Add support for models.create endpoint

Signed-off-by: Mattt Zmuda <mattt@replicate.com>

Author: mattt

replicate/model.py

@@ -150,6 +150,63 @@ def get(self, key: str) -> Model:
         resp = self._client._request("GET", f"/v1/models/{key}")
         return self._prepare_model(resp.json())
 
+    def create(  # pylint: disable=arguments-differ disable=too-many-arguments
+        self,
+        owner: str,
+        name: str,
+        *,
+        visibility: str,
+        hardware: str,
+        description: Optional[str] = None,
+        github_url: Optional[str] = None,
+        paper_url: Optional[str] = None,
+        license_url: Optional[str] = None,
+        cover_image_url: Optional[str] = None,
+    ) -> Model:
+        """
+        Create a model.
+
+        Args:
+            owner: The name of the user or organization that will own the model.
+            name: The name of the model.
+            visibility: Whether the model should be public or private.
+            hardware: The SKU for the hardware used to run the model. Possible values can be found by calling `replicate.hardware.list()`.
+            description: A description of the model.
+            github_url: A URL for the model's source code on GitHub.
+            paper_url: A URL for the model's paper.
+            license_url: A URL for the model's license.
+            cover_image_url: A URL for the model's cover image.
+
+        Returns:
+            The created model.
+        """
+
+        body = {
+            "owner": owner,
+            "name": name,
+            "visibility": visibility,
+            "hardware": hardware,
+        }
+
+        if description is not None:
+            body["description"] = description
+
+        if github_url is not None:
+            body["github_url"] = github_url
+
+        if paper_url is not None:
+            body["paper_url"] = paper_url
+
+        if license_url is not None:
+            body["license_url"] = license_url
+
+        if cover_image_url is not None:
+            body["cover_image_url"] = cover_image_url
+
+        resp = self._client._request("POST", "/v1/models", json=body)
+
+        return self._prepare_model(resp.json())
+
     def _prepare_model(self, attrs: Union[Model, Dict]) -> Model:
         if isinstance(attrs, BaseModel):
             attrs.id = f"{attrs.owner}/{attrs.name}"

tests/cassettes/models-create.yaml

@@ -0,0 +1,80 @@
+interactions:
+- request:
+    body: '{"owner": "test", "name": "python-example", "visibility": "private", "hardware":
+      "cpu", "description": "An example model"}'
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '123'
+      content-type:
+      - application/json
+      host:
+      - api.replicate.com
+      user-agent:
+      - replicate-python/0.15.6
+    method: POST
+    uri: https://api.replicate.com/v1/models
+  response:
+    content: '{"url": "https://replicate.com/test/python-example", "owner": "test",
+      "name": "python-example", "description": "An example model", "visibility": "private",
+      "github_url": null, "paper_url": null, "license_url": null, "run_count": 0,
+      "cover_image_url": null, "default_example": null, "latest_version": null}'
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 81ff2e098ec0eb5b-SEA
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '307'
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 02 Nov 2023 20:38:12 GMT
+      Server:
+      - cloudflare
+      Strict-Transport-Security:
+      - max-age=15552000
+      allow:
+      - GET, POST, HEAD, OPTIONS
+      content-security-policy-report-only:
+      - 'font-src ''report-sample'' ''self'' data:; img-src ''report-sample'' ''self''
+        data: https://replicate.delivery https://*.replicate.delivery https://*.githubusercontent.com
+        https://github.com; script-src ''report-sample'' ''self'' https://cdn.rudderlabs.com/v1.1/rudder-analytics.min.js;
+        style-src ''report-sample'' ''self'' ''unsafe-inline''; connect-src ''report-sample''
+        ''self'' https://replicate.delivery https://*.replicate.delivery https://*.rudderlabs.com
+        https://*.rudderstack.com https://*.mux.com https://*.sentry.io; worker-src
+        ''none''; media-src ''report-sample'' ''self'' https://replicate.delivery
+        https://*.replicate.delivery https://*.mux.com https://*.sentry.io; default-src
+        ''self''; report-uri'
+      cross-origin-opener-policy:
+      - same-origin
+      nel:
+      - '{"report_to":"heroku-nel","max_age":3600,"success_fraction":0.005,"failure_fraction":0.05,"response_headers":["Via"]}'
+      ratelimit-remaining:
+      - '2999'
+      ratelimit-reset:
+      - '1'
+      referrer-policy:
+      - same-origin
+      report-to:
+      - '{"group":"heroku-nel","max_age":3600,"endpoints":[{"url":"https://nel.heroku.com/reports?ts=1698957492&sid=1b10b0ff-8a76-4548-befa-353fc6c6c045&s=m%2Fs583uNWdN4J4bm1G3JZoilUVMbh89egg%2FAEcTPZm4%3D"}]}'
+      reporting-endpoints:
+      - heroku-nel=https://nel.heroku.com/reports?ts=1698957492&sid=1b10b0ff-8a76-4548-befa-353fc6c6c045&s=m%2Fs583uNWdN4J4bm1G3JZoilUVMbh89egg%2FAEcTPZm4%3D
+      vary:
+      - Cookie, origin
+      via:
+      - 1.1 vegur, 1.1 google
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - DENY
+    http_version: HTTP/1.1
+    status_code: 201
+version: 1

tests/test_hardware.py

@@ -1,6 +1,4 @@
-import httpx
 import pytest
-import respx
 
 import replicate
 

tests/test_model.py

@@ -27,3 +27,34 @@ async def test_models_list(mock_replicate_api_token):
     assert models[0].owner is not None
     assert models[0].name is not None
     assert models[0].visibility == "public"
+
+
+@pytest.mark.vcr("models-create.yaml")
+@pytest.mark.asyncio
+async def test_models_create(mock_replicate_api_token):
+    model = replicate.models.create(
+        owner="test",
+        name="python-example",
+        visibility="private",
+        hardware="cpu",
+        description="An example model",
+    )
+
+    assert model.owner == "test"
+    assert model.name == "python-example"
+    assert model.visibility == "private"
+
+
+@pytest.mark.vcr("models-create.yaml")
+@pytest.mark.asyncio
+async def test_models_create_with_positional_arguments(mock_replicate_api_token):
+    model = replicate.models.create(
+        "test",
+        "python-example",
+        visibility="private",
+        hardware="cpu",
+    )
+
+    assert model.owner == "test"
+    assert model.name == "python-example"
+    assert model.visibility == "private"