Add support for models.search endpoint

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

Author: mattt

replicate/model.py

@@ -207,6 +207,46 @@ async def async_list(
 
         return Page[Model](**obj)
 
+    def search(self, query: str) -> Page[Model]:
+        """
+        Search for public models.
+
+        Parameters:
+            query: The search query.
+        Returns:
+            Page[Model]: A page of models matching the search query.
+        """
+        resp = self._client._request(
+            "QUERY", "/v1/models", content=query, headers={"Content-Type": "text/plain"}
+        )
+
+        obj = resp.json()
+        obj["results"] = [
+            _json_to_model(self._client, result) for result in obj["results"]
+        ]
+
+        return Page[Model](**obj)
+
+    async def async_search(self, query: str) -> Page[Model]:
+        """
+        Asynchronously search for public models.
+
+        Parameters:
+            query: The search query.
+        Returns:
+            Page[Model]: A page of models matching the search query.
+        """
+        resp = await self._client._async_request(
+            "QUERY", "/v1/models", content=query, headers={"Content-Type": "text/plain"}
+        )
+
+        obj = resp.json()
+        obj["results"] = [
+            _json_to_model(self._client, result) for result in obj["results"]
+        ]
+
+        return Page[Model](**obj)
+
     @overload
     def get(self, key: str) -> Model: ...
 

tests/test_model.py

@@ -1,6 +1,7 @@
 import pytest
 
 import replicate
+from replicate.model import Model, Page
 
 
 @pytest.mark.vcr("models-get.yaml")
@@ -130,3 +131,26 @@ async def test_models_predictions_create(async_flag):
     # assert prediction.model == "meta/llama-2-70b-chat"
     assert prediction.model == "replicate/lifeboat-70b"  # FIXME: this is temporary
     assert prediction.status == "starting"
+
+
+@pytest.mark.vcr("models-search.yaml")
+@pytest.mark.asyncio
+@pytest.mark.parametrize("async_flag", [True, False])
+async def test_models_search(async_flag):
+    query = "llama"
+
+    if async_flag:
+        page = await replicate.models.async_search(query)
+    else:
+        page = replicate.models.search(query)
+
+    assert isinstance(page, Page)
+    assert len(page.results) > 0
+
+    for model in page.results:
+        assert isinstance(model, Model)
+        assert model.id is not None
+        assert model.owner is not None
+        assert model.name is not None
+
+    assert any("meta" in model.name.lower() for model in page.results)