Make openapi_schema consistently a method

This commit also fixes the implementation to return the OpenAPI schema
for the requested version rather than always the latest version.

Author: aron

replicate/use.py

@@ -5,7 +5,6 @@
 import hashlib
 import os
 import tempfile
-from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
 from typing import (
@@ -25,7 +24,6 @@
     cast,
     overload,
 )
-from urllib.parse import urlparse
 
 import httpx
 
@@ -62,36 +60,6 @@ def _has_concatenate_iterator_output_type(openapi_schema: dict) -> bool:
     return True
 
 
-def _has_iterator_output_type(openapi_schema: dict) -> bool:
-    """
-    Returns true if the model output type is an iterator (non-concatenate).
-    """
-    output = openapi_schema.get("components", {}).get("schemas", {}).get("Output", {})
-    return (
-        output.get("type") == "array" and output.get("x-cog-array-type") == "iterator"
-    )
-
-
-def _download_file(url: str) -> Path:
-    """
-    Download a file from URL to a temporary location and return the Path.
-    """
-    parsed_url = urlparse(url)
-    filename = os.path.basename(parsed_url.path)
-
-    if not filename or "." not in filename:
-        filename = "download"
-
-    _, ext = os.path.splitext(filename)
-    with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_file:
-        with httpx.stream("GET", url) as response:
-            response.raise_for_status()
-            for chunk in response.iter_bytes():
-                temp_file.write(chunk)
-
-    return Path(temp_file.name)
-
-
 def _process_iterator_item(item: Any, openapi_schema: dict) -> Any:
     """
     Process a single item from an iterator output based on schema.
@@ -357,7 +325,6 @@ class FunctionRef(Protocol, Generic[Input, Output]):
     __call__: Callable[Input, Output]
 
 
-@dataclass
 class Run[O]:
     """
     Represents a running prediction with access to the underlying schema.
@@ -416,13 +383,13 @@ def logs(self) -> Optional[str]:
         return self._prediction.logs
 
 
-@dataclass
 class Function(Generic[Input, Output]):
     """
     A wrapper for a Replicate model that can be called as a function.
     """
 
     _ref: str
+    _streaming: bool
 
     def __init__(self, ref: str, *, streaming: bool) -> None:
         self._ref = ref
@@ -460,7 +427,9 @@ def create(self, *_: Input.args, **inputs: Input.kwargs) -> Run[Output]:
             )
 
         return Run(
-            prediction=prediction, schema=self.openapi_schema, streaming=self._streaming
+            prediction=prediction,
+            schema=self.openapi_schema(),
+            streaming=self._streaming,
         )
 
     @property
@@ -470,18 +439,26 @@ def default_example(self) -> Optional[dict[str, Any]]:
         """
         raise NotImplementedError("This property has not yet been implemented")
 
-    @cached_property
     def openapi_schema(self) -> dict[str, Any]:
         """
         Get the OpenAPI schema for this model version.
         """
-        latest_version = self._model.latest_version
-        if latest_version is None:
+        return self._openapi_schema
+
+    @cached_property
+    def _openapi_schema(self) -> dict[str, Any]:
+        _, _, model_version = self._parsed_ref
+        model = self._model
+
+        version = (
+            model.versions.get(model_version) if model_version else model.latest_version
+        )
+        if version is None:
             msg = f"Model {self._model.owner}/{self._model.name} has no latest version"
             raise ValueError(msg)
 
-        schema = latest_version.openapi_schema
-        if cog_version := latest_version.cog_version:
+        schema = version.openapi_schema
+        if cog_version := version.cog_version:
             schema = make_schema_backwards_compatible(schema, cog_version)
         return _dereference_schema(schema)
 
@@ -524,7 +501,6 @@ def _version(self) -> Version | None:
         return version
 
 
-@dataclass
 class AsyncRun[O]:
     """
     Represents a running prediction with access to its version (async version).
@@ -583,21 +559,25 @@ async def logs(self) -> Optional[str]:
         return self._prediction.logs
 
 
-@dataclass
 class AsyncFunction(Generic[Input, Output]):
     """
     An async wrapper for a Replicate model that can be called as a function.
     """
 
-    function_ref: str
-    streaming: bool
+    _ref: str
+    _streaming: bool
+    _openapi_schema: dict[str, Any] | None = None
+
+    def __init__(self, ref: str, *, streaming: bool) -> None:
+        self._ref = ref
+        self._streaming = streaming
 
     def _client(self) -> Client:
         return Client()
 
     @cached_property
     def _parsed_ref(self) -> Tuple[str, str, Optional[str]]:
-        return ModelVersionIdentifier.parse(self.function_ref)
+        return ModelVersionIdentifier.parse(self._ref)
 
     async def _model(self) -> Model:
         client = self._client()
@@ -662,7 +642,7 @@ async def create(self, *_: Input.args, **inputs: Input.kwargs) -> AsyncRun[Outpu
         return AsyncRun(
             prediction=prediction,
             schema=await self.openapi_schema(),
-            streaming=self.streaming,
+            streaming=self._streaming,
         )
 
     @property
@@ -676,16 +656,26 @@ async def openapi_schema(self) -> dict[str, Any]:
         """
         Get the OpenAPI schema for this model version asynchronously.
         """
-        model = await self._model()
-        latest_version = model.latest_version
-        if latest_version is None:
-            msg = f"Model {model.owner}/{model.name} has no latest version"
-            raise ValueError(msg)
+        if not self._openapi_schema:
+            _, _, model_version = self._parsed_ref
 
-        schema = latest_version.openapi_schema
-        if cog_version := latest_version.cog_version:
-            schema = make_schema_backwards_compatible(schema, cog_version)
-        return _dereference_schema(schema)
+            model = await self._model()
+            if model_version:
+                version = await model.versions.async_get(model_version)
+            else:
+                version = model.latest_version
+
+            if version is None:
+                msg = f"Model {model.owner}/{model.name} has no version"
+                raise ValueError(msg)
+
+            schema = version.openapi_schema
+            if cog_version := version.cog_version:
+                schema = make_schema_backwards_compatible(schema, cog_version)
+
+            self._openapi_schema = _dereference_schema(schema)
+
+        return self._openapi_schema
 
 
 @overload

tests/test_use.py

@@ -372,7 +372,7 @@ async def test_use_function_openapi_schema_dereferenced(client_mode):
     if client_mode == ClientMode.ASYNC:
         schema = await hotdog_detector.openapi_schema()
     else:
-        schema = hotdog_detector.openapi_schema
+        schema = hotdog_detector.openapi_schema()
 
     assert schema["components"]["schemas"]["Output"] == {
         "type": "object",