Add overloads to models.get method

Signed-off-by: Mattt Zmuda <[email protected]>

Author: mattt

replicate/model.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Union, overload
 
 from typing_extensions import NotRequired, TypedDict, Unpack, deprecated
 
@@ -207,31 +207,65 @@ async def async_list(
 
         return Page[Model](**obj)
 
+    @overload
     def get(self, key: str) -> Model:
         """
         Get a model by name.
 
         Args:
-            key: The qualified name of the model, in the format `owner/model-name`.
+            key: The qualified name of the model, in the format `owner/name`.
         Returns:
             The model.
         """
+        ...
 
-        resp = self._client._request("GET", f"/v1/models/{key}")
+    @overload
+    def get(self, owner: str, name: str) -> Model:
+        """
+        Get a model by owner and name.
+
+        Args:
+            owner: The owner of the model.
+            name: The name of the model.
+        Returns:
+            The model.
+        """
+        ...
+
+    def get(self, *args, **kwargs) -> Model:
+        url = _get_model_url(*args, **kwargs)
+        resp = self._client._request("GET", url)
 
         return _json_to_model(self._client, resp.json())
 
+    @overload
     async def async_get(self, key: str) -> Model:
         """
         Get a model by name.
 
         Args:
-            key: The qualified name of the model, in the format `owner/model-name`.
+            key: The qualified name of the model, in the format `owner/name`.
         Returns:
             The model.
         """
+        ...
 
-        resp = await self._client._async_request("GET", f"/v1/models/{key}")
+    @overload
+    async def async_get(self, owner: str, name: str) -> Model:
+        """
+        Get a model by owner and name.
+
+        Args:
+            owner: The owner of the model.
+            name: The name of the model.
+        Returns:
+            The model.
+        """
+        ...
+
+    async def async_get(self, *args, **kwargs) -> Model:
+        url = _get_model_url(*args, **kwargs)
+        resp = await self._client._async_request("GET", url)
 
         return _json_to_model(self._client, resp.json())
 
@@ -288,6 +322,13 @@ async def async_create(
 
         return _json_to_model(self._client, resp.json())
 
+    def delete(self, key: str) -> None:
+        """
+        Delete a model.
+        """
+
+        self._client._request("DELETE", f"/v1/models/{key}")
+
 
 class ModelsPredictions(Namespace):
     """
@@ -374,6 +415,50 @@ def _create_model_body(  # pylint: disable=too-many-arguments
     return body
 
 
+def _get_model_url(*args, **kwargs) -> str:
+    """
+    Extracts the key from the provided arguments.
+
+    Args:
+        args: Positional arguments that may contain the key or owner and name.
+        kwargs: Keyword arguments that may contain the key, owner, or name.
+
+    Returns:
+        The extracted key in the format 'owner/name'.
+
+    Raises:
+        ValueError: If the arguments are invalid or insufficient to determine the key.
+    """
+    if len(args) > 0 and len(kwargs) > 0:
+        raise ValueError("Cannot mix positional and keyword arguments")
+
+    owner = kwargs.get("owner", None)
+    name = kwargs.get("name", None)
+    key = kwargs.get("key", None)
+
+    if key and (owner or name):
+        raise ValueError(
+            "Must specify exactly one of 'owner' and 'name' or single 'key' in the format 'owner/name'"
+        )
+
+    if args:
+        if len(args) == 1:
+            key = args[0]
+        elif len(args) == 2:
+            owner, name = args
+        else:
+            raise ValueError("Invalid number of arguments")
+
+    if not key:
+        if not (owner and name):
+            raise ValueError(
+                "Both 'owner' and 'name' must be provided if 'key' is not specified."
+            )
+        key = f"{owner}/{name}"
+
+    return f"/v1/models/{key}"
+
+
 def _json_to_model(client: "Client", json: Dict[str, Any]) -> Model:
     model = Model(**json)
     model._client = client