Add some docstrings

great_ai/deploy/great_ai.py

@@ -49,6 +49,8 @@ class GreatAI(Generic[T, V]):
         self,
         func: Callable[..., Union[V, Awaitable[V]]],
     ):
+        """Do not call this function directly, use GreatAI.create instead."""
+
         func = automatically_decorate_parameters(func)
         get_function_metadata_store(func).is_finalised = True
 
@@ -94,6 +96,43 @@ class GreatAI(Generic[T, V]):
     def create(
         func: Union[Callable[..., Awaitable[V]], Callable[..., V]],
     ) -> Union["GreatAI[Awaitable[Trace[V]], V]", "GreatAI[Trace[V], V]"]:
+        """Decorate a function by wrapping it in a GreatAI instance.
+
+        The function can be typed, synchronous or async. If it has
+        unwrapped parameters (parameters not affected by a @parameter
+        or @use_model decorator), those will be automatically wrapped.
+
+        The return value is replaced by a Trace (or Awaitable[Trace]),
+        while the original return value is available under the `.output`
+        property.
+
+        For configuration options, see great_ai.configure.
+
+        Examples:
+            >>> @GreatAI.create
+            ... def my_function(a):
+            ...     return a + 2
+            >>> my_function(3).output
+            5
+
+            >>> @GreatAI.create
+            ... def my_function(a: int) -> int:
+            ...     return a + 2
+            >>> my_function(3)
+            Trace[int]...
+
+            >>> my_function('3').output
+            Traceback (most recent call last):
+                ...
+            TypeError: type of a must be int; got str instead
+
+        Args:
+            func: The prediction function that needs to be decorated.
+
+        Returns:
+            A GreatAI instance wrapping `func`.
+
+        """
         return GreatAI[Trace[V], V](
             func,
         )

great_ai/models/save_model.py

@@ -9,6 +9,30 @@ from ..context import get_context
 def save_model(
     model: Union[Path, str, object], key: str, *, keep_last_n: Optional[int] = None
 ) -> str:
+    """Save (and optionally serialise) a model in order to use by `use_model`.
+
+    The `model` can be a Path or string representing a path in which case the
+    local file/folder is read and saved using the current LargeFile implementation.
+    In case `model` is an object, it is serialised using `dill` before uploading it.
+
+    Examples:
+            >>> from great_ai import use_model
+            >>> save_model(3, 'my_number')
+            'my_number:...'
+            
+            >>> @use_model('my_number')
+            ... def my_function(a, model):
+            ...     return a + model
+            >>> my_function(4)
+            7
+
+    Args:
+        model: The object or path to be uploaded.
+        key: The model's name.
+        keep_last_n: If specified, remove old models and only keep the latest n. Directly passed to LargeFile.
+    Returns:
+        The key and version of the saved model separated by a colon. Example: "key:version"
+    """
     file = get_context().large_file_implementation(
         name=key, mode="wb", keep_last_n=keep_last_n
     )

great_ai/models/use_model.py

@@ -30,6 +30,38 @@ def use_model(
     version: Union[int, Literal["latest"]] = "latest",
     model_kwarg_name: str = "model",
 ) -> Callable[[F], F]:
+    """Inject a model into a function.
+
+    Load a model specified by `key` and `version` using the
+    currently active `LargeFile` implementation. If it's a
+    single object, it is deserialised using `dill`. If it's
+    a directory of files, a `pathlib.Path` instance is given.
+
+    By default, the function's `model` parameter is replaced
+    by the loaded model. This can be customised by changing
+    `model_kwarg_name`.
+
+    Multiple models can be loaded by decorating the same
+    function with `use_model` multiple times.
+
+    Examples:
+            >>> from great_ai import save_model
+            >>> save_model(3, 'my_number')
+            'my_number:...'
+            >>> @use_model('my_number')
+            ... def my_function(a, model):
+            ...     return a + model
+            >>> my_function(4)
+            7
+
+        Args:
+            key: The model's name as stored by the LargeFile implementation.
+            version: The model's version as stored by the LargeFile implementation.
+            model_kwarg_name: the parameter to use for injecting the loaded model
+        Returns:
+            A decorator for model injection.
+    """
+
     assert (
         isinstance(version, int) or version == "latest"
     ), "Only integers or the string literal `latest` is allowed as a version"

great_ai/parameters/parameter.py

@@ -17,6 +17,35 @@ def parameter(
     validator: Callable[[Any], bool] = lambda _: True,
     disable_logging: bool = False,
 ) -> Callable[[F], F]:
+    """Control the validation and logging of function parameters.
+
+    Examples:
+        >>> @parameter('a')
+        ... def my_function(a: int):
+        ...     return a + 2
+        >>> my_function(4)
+        6
+        >>> my_function('3')
+        Traceback (most recent call last):
+            ...
+        TypeError: type of a must be int; got str instead
+
+        >>> @parameter('positive_a', validator=lambda v: v > 0)
+        ... def my_function(positive_a: int):
+        ...     return a + 2
+        >>> my_function(-1)
+        Traceback (most recent call last):
+            ...
+        great_ai.errors.argument_validation_error.ArgumentValidationError: ...
+
+    Args:
+        parameter_name: Name of parameter to consider
+        validator: Optional validator to run against the concrete argument. ArgumentValidationError is thrown when the return value is False.
+        disable_logging: Do not save the value in any active TracingContext.
+    Returns:
+        A decorator for argument validation.
+    """
+
     def decorator(func: F) -> F:
         get_function_metadata_store(func).input_parameter_names.append(parameter_name)
         assert_function_is_not_finalised(func)

great_ai/utilities/unique.py

@@ -11,8 +11,10 @@ def unique(values: Iterable[T], *, key: Callable[[T], Any] = lambda v: v) -> Lis
     Examples:
         >>> unique([1, 1, 5, 3, 3])
         [1, 5, 3]
+
         >>> unique([{'a': 1, 'b': 2}, {'a': 1, 'b': 3}], key=lambda v: v['a'])
         [{'a': 1, 'b': 2}]
+
         >>> unique([{'a': 1, 'b': 2}, {'a': 1, 'b': 3}], key=lambda v: v['b'])
         [{'a': 1, 'b': 2}, {'a': 1, 'b': 3}]