Add support for typing use() function

Author: aron

README.md

@@ -592,16 +592,49 @@ for output in outputs:
     print(get_url_path(output)) # "https://replicate.delivery/xyz"
 ```
 
+### Typing
+
+By default `use()` knows nothing about the interface of the model. To provide a better developer experience we provide two methods to add type annotations to the function returned by the `use()` helper.
+
+**1. Provide a function signature**
+
+The use method accepts a function signature as an additional `hint` keyword argument. When provided it will use this signature for the `model()` and `model.create()` functions.
+
+```py
+# Flux takes a required prompt string and optional image and seed.
+def hint(*, prompt: str, image: Path | None = None, seed: int | None = None) -> str: ...
+
+flux_dev = use("black-forest-labs/flux-dev", hint=hint)
+output1 = flux_dev() # will warn that `prompt` is missing
+output2 = flux_dev(prompt="str") # output2 will be typed as `str`
+```
+
+**2. Provide a class**
+
+The second method requires creating a callable class with a `name` field. The name will be used as the function reference when passed to `use()`.
+
+```py
+class FluxDev:
+    name = "black-forest-labs/flux-dev"
+
+    def __call__( self, *, prompt: str, image: Path | None = None, seed: int | None = None ) -> str: ...
+
+flux_dev = use(FluxDev)
+output1 = flux_dev() # will warn that `prompt` is missing
+output2 = flux_dev(prompt="str") # output2 will be typed as `str`
+```
+
+In future we hope to provide tooling to generate and provide these models as packages to make working with them easier. For now you may wish to create your own.
+
 ### TODO
 
 There are several key things still outstanding:
 
  1. Support for asyncio.
- 2. Support for typing the return value.
- 3. Support for streaming text when available (rather than polling)
- 4. Support for streaming files when available (rather than polling)
- 5. Support for cleaning up downloaded files.
- 6. Support for streaming logs using `OutputIterator`.
+ 2. Support for streaming text when available (rather than polling)
+ 3. Support for streaming files when available (rather than polling)
+ 4. Support for cleaning up downloaded files.
+ 5. Support for streaming logs using `OutputIterator`.
 
 ## Development
 

replicate/use.py

@@ -4,11 +4,24 @@
 # - [ ] Support asyncio variant
 import inspect
 import os
+import sys
 import tempfile
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
-from typing import Any, Iterator, Optional, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterator,
+    Optional,
+    ParamSpec,
+    Protocol,
+    Tuple,
+    TypeVar,
+    cast,
+    overload,
+)
 from urllib.parse import urlparse
 
 import httpx
@@ -24,16 +37,35 @@
 __all__ = ["use", "get_path_url"]
 
 
+def _in_repl() -> bool:
+    return bool(
+        sys.flags.interactive  # python -i
+        or hasattr(sys, "ps1")  # prompt strings exist
+        or (
+            sys.stdin.isatty()  # tty
+            and sys.stdout.isatty()
+        )
+        or ("get_ipython" in globals())
+    )
+
+
 def _in_module_scope() -> bool:
     """
     Returns True when called from top level module scope.
     """
     if os.getenv("REPLICATE_ALWAYS_ALLOW_USE"):
         return True
 
+    # If we're running in a REPL.
+    if _in_repl():
+        return True
+
     if frame := inspect.currentframe():
+        print(frame)
         if caller := frame.f_back:
+            print(caller.f_code.co_name)
             return caller.f_code.co_name == "<module>"
+
     return False
 
 
@@ -248,16 +280,26 @@ def get_path_url(path: Any) -> str | None:
         return None
 
 
+Input = ParamSpec("Input")
+Output = TypeVar("Output")
+
+
+class FunctionRef(Protocol, Generic[Input, Output]):
+    name: str
+
+    __call__: Callable[Input, Output]
+
+
 @dataclass
-class Run:
+class Run[O]:
     """
     Represents a running prediction with access to its version.
     """
 
     prediction: Prediction
     schema: dict
 
-    def output(self) -> Union[Any, Iterator[Any]]:
+    def output(self) -> O:
         """
         Wait for the prediction to complete and return its output.
         """
@@ -269,10 +311,13 @@ def output(self) -> Union[Any, Iterator[Any]]:
         # Return an OutputIterator for iterator output types (including concatenate iterators)
         if _has_iterator_output_type(self.schema):
             is_concatenate = _has_concatenate_iterator_output_type(self.schema)
-            return OutputIterator(
-                lambda: self.prediction.output_iterator(),
-                self.schema,
-                is_concatenate=is_concatenate,
+            return cast(
+                O,
+                OutputIterator(
+                    lambda: self.prediction.output_iterator(),
+                    self.schema,
+                    is_concatenate=is_concatenate,
+                ),
             )
 
         # Process output for file downloads based on schema
@@ -288,7 +333,7 @@ def logs(self) -> Optional[str]:
 
 
 @dataclass
-class Function:
+class Function(Generic[Input, Output]):
     """
     A wrapper for a Replicate model that can be called as a function.
     """
@@ -333,11 +378,10 @@ def _version(self) -> Version | None:
 
         return version
 
-    def __call__(self, **inputs: Any) -> Any:
-        run = self.create(**inputs)
-        return run.output()
+    def __call__(self, *args: Input.args, **inputs: Input.kwargs) -> Output:
+        return self.create(*args, **inputs).output()
 
-    def create(self, **inputs: Any) -> Run:
+    def create(self, *_: Input.args, **inputs: Input.kwargs) -> Run[Output]:
         """
         Start a prediction with the specified inputs.
         """
@@ -365,14 +409,14 @@ def create(self, **inputs: Any) -> Run:
         return Run(prediction, self.openapi_schema)
 
     @property
-    def default_example(self) -> Optional[Prediction]:
+    def default_example(self) -> Optional[dict[str, Any]]:
         """
         Get the default example for this model.
         """
         raise NotImplementedError("This property has not yet been implemented")
 
     @cached_property
-    def openapi_schema(self) -> dict[Any, Any]:
+    def openapi_schema(self) -> dict[str, Any]:
         """
         Get the OpenAPI schema for this model version.
         """
@@ -387,7 +431,19 @@ def openapi_schema(self) -> dict[Any, Any]:
         return schema
 
 
-def use(function_ref: str) -> Function:
+@overload
+def use(ref: FunctionRef[Input, Output]) -> Function[Input, Output]: ...
+
+
+@overload
+def use(ref: str, *, hint: Callable[Input, Output]) -> Function[Input, Output]: ...
+
+
+def use(
+    ref: str | FunctionRef[Input, Output],
+    *,
+    hint: Callable[Input, Output] | None = None,
+) -> Function[Input, Output]:
     """
     Use a Replicate model as a function.
 
@@ -402,4 +458,9 @@ def use(function_ref: str) -> Function:
     if not _in_module_scope():
         raise RuntimeError("You may only call replicate.use() at the top level.")
 
-    return Function(function_ref)
+    try:
+        ref = ref.name  # type: ignore
+    except AttributeError:
+        pass
+
+    return Function(str(ref))

tests/test_use.py

@@ -280,6 +280,27 @@ async def test_use_with_version_identifier(client_mode):
     assert output == "not hotdog"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("client_mode", [ClientMode.DEFAULT])
+@respx.mock
+async def test_use_with_function_ref(client_mode):
+    mock_model_endpoints()
+    mock_prediction_endpoints()
+
+    class HotdogDetector:
+        name = "acme/hotdog-detector:xyz123"
+
+        def __call__(self, prompt: str) -> str: ...
+
+    hotdog_detector = replicate.use(HotdogDetector())
+
+    # Call function with prompt="hello world"
+    output = hotdog_detector(prompt="hello world")
+
+    # Assert that output is the completed output from the prediction request
+    assert output == "not hotdog"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("client_mode", [ClientMode.DEFAULT])
 @respx.mock