Skip to content

Commit d69f588

Browse files
committed
Add overloads to models.delete method
Signed-off-by: Mattt Zmuda <[email protected]>
1 parent db719ef commit d69f588

File tree

1 file changed

+41
-7
lines changed

1 file changed

+41
-7
lines changed

replicate/model.py

Lines changed: 41 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,43 @@ async def async_get(self, *args, **kwargs) -> Model:
244244

245245
return _json_to_model(self._client, resp.json())
246246

247+
@overload
248+
def delete(self, key: str) -> Model: ...
249+
250+
@overload
251+
def delete(self, owner: str, name: str) -> Model: ...
252+
253+
def delete(self, *args, **kwargs) -> Model:
254+
"""
255+
Delete a model by name.
256+
"""
257+
258+
url = _delete_model_url(*args, **kwargs)
259+
resp = self._client._request("DELETE", url)
260+
261+
return _json_to_model(self._client, resp.json())
262+
263+
@overload
264+
async def async_delete(self, key: str) -> Model: ...
265+
266+
@overload
267+
async def async_delete(self, owner: str, name: str) -> Model: ...
268+
269+
async def async_delete(self, *args, **kwargs) -> Model:
270+
"""
271+
Delete a model by name.
272+
273+
Args:
274+
key: The qualified name of the model, in the format `owner/name`.
275+
Returns:
276+
The model.
277+
"""
278+
279+
url = _delete_model_url(*args, **kwargs)
280+
resp = await self._client._async_request("DELETE", url)
281+
282+
return _json_to_model(self._client, resp.json())
283+
247284
class CreateModelParams(TypedDict):
248285
"""Parameters for creating a model."""
249286

@@ -297,13 +334,6 @@ async def async_create(
297334

298335
return _json_to_model(self._client, resp.json())
299336

300-
def delete(self, key: str) -> None:
301-
"""
302-
Delete a model.
303-
"""
304-
305-
self._client._request("DELETE", f"/v1/models/{key}")
306-
307337

308338
class ModelsPredictions(Namespace):
309339
"""
@@ -421,6 +451,10 @@ def _get_model_url(*args, **kwargs) -> str:
421451
return f"/v1/models/{key}"
422452

423453

454+
def _delete_model_url(*args, **kwargs) -> str:
455+
return _get_model_url(*args, **kwargs)
456+
457+
424458
def _json_to_model(client: "Client", json: Dict[str, Any]) -> Model:
425459
model = Model(**json)
426460
model._client = client

0 commit comments

Comments
 (0)