dyff-schema 0.10.3__py3-none-any.whl → 0.12.0__py3-none-any.whl

dyff/schema/annotations.py

@@ -0,0 +1,297 @@
+# SPDX-FileCopyrightText: 2024 UL Research Institutes
+# SPDX-License-Identifier: Apache-2.0
+
+# mypy: disable-error-code="import-untyped"
+import functools
+import typing
+from typing import Callable, Generic, Literal, NamedTuple, Optional, TypeVar
+
+import pyarrow.dataset
+import pydantic
+from typing_extensions import ParamSpec
+
+from dyff.schema.dataset import ReplicatedItem, arrow
+from dyff.schema.platform import (
+    DataSchema,
+    Dataset,
+    Evaluation,
+    Measurement,
+    MeasurementLevel,
+    MeasurementSpec,
+    MethodImplementation,
+    MethodImplementationKind,
+    MethodImplementationPythonFunction,
+    MethodInput,
+    MethodInputKind,
+    MethodOutput,
+    MethodOutputKind,
+    MethodParameter,
+    MethodScope,
+)
+from dyff.schema.requests import MethodCreateRequest
+
+
+def _fqn(obj) -> tuple[str, str]:
+    """See: https://stackoverflow.com/a/70693158"""
+    try:
+        module = obj.__module__
+    except AttributeError:
+        module = obj.__class__.__module__
+    try:
+        name = obj.__qualname__
+    except AttributeError:
+        name = obj.__class__.__qualname__
+    # if obj is a method of builtin class, then module will be None
+    if module == "builtins" or module is None:
+        raise AssertionError("should not be called on a builtin")
+    return module, name
+
+
+class DataAnnotation(NamedTuple):
+    kind: str
+    direction: Literal["input", "output"]
+    level: Optional[MeasurementLevel] = None
+    schema: Optional[DataSchema] = None
+
+
+def Input(input_type) -> DataAnnotation:
+    """Apply this annotation to parameters of a Method implementation to
+    indicate that the parameter expects a PyArrow dataset derived from the
+    specified type of entity, e.g.::
+
+        def my_method(input_data: Annotated[pyarrow.dataset.Dataset, Input(Evaluation)], ...
+
+    :param input_type: A Dyff entity type with associated input data; one of
+        {Dataset, Evaluation, Measurement}
+    :return: Annotation data
+    """
+    if input_type == Dataset:
+        return DataAnnotation(kind="Dataset", direction="input")
+    elif input_type == Evaluation:
+        return DataAnnotation(kind="Evaluation", direction="input")
+    elif input_type == Measurement:
+        return DataAnnotation(kind="Measurement", direction="input")
+    else:
+        raise TypeError()
+
+
+# TODO: I think this could work if we ever upgrade to Python 3.12+. We need the
+# type checker to accept `InputData[Evaluation]` and treat it as an alias for
+# `pyarrow.dataset.Dataset`.
+#
+# if typing.TYPE_CHECKING:
+#     _InputDataType = TypeVar("_InputDataType")
+#     type InputData[_InputDataType] = pyarrow.dataset.Dataset
+# else:
+#
+#     class InputData:
+#         def __init__(self):
+#             raise NotImplementedError()
+#
+#         def __class_getitem__(cls, input_type) -> typing.GenericAlias:
+#             return Annotated[pyarrow.dataset.Dataset, Input(input_type)]
+
+
+def Output(output_type, *, schema, level: Optional[MeasurementLevel] = None):
+    """Apply this annotation to the return type of a Method to provide
+    metadata about the type of output created by the Method, e.g.::
+
+        def my_method(...) -> Annotated[
+            Iterable[pyarrow.RecordBatch],
+            Output(Measurement, schema=MyPydanticType, level=MeasurementLevel.Instance)
+        ]: ...
+
+    :param output_type: A Dyff entity type with associated output data; one of
+        {Measurement, SafetyCase}
+    :param schema: The schema of the output. Can be a type derived from
+        pydantic.BaseModel or an Arrow schema. The mandatory fields `_index_`
+        and `_replication_` will be *added* and should not be present.
+    :param level: The MeasurementLevel, if the output is a Measurement.
+    :return: Annotation data
+    """
+    if isinstance(schema, type) and issubclass(schema, pydantic.BaseModel):
+        RowSchema = pydantic.create_model(
+            "RowSchema", __base__=(schema, ReplicatedItem)
+        )
+        data_schema = DataSchema(
+            arrowSchema=arrow.encode_schema(arrow.arrow_schema(RowSchema))
+        )
+    elif isinstance(schema, pyarrow.Schema):
+        raise NotImplementedError()
+        # TODO: Add _index_ and _replication_
+        # data_schema = DataSchema(arrowSchema=arrow.encode_schema(schema))
+    else:
+        raise TypeError()
+
+    if output_type == Measurement:
+        if level is None:
+            raise ValueError("Must specify 'level' when output_type == Measurement")
+        return DataAnnotation(
+            kind="Measurement",
+            direction="output",
+            level=level,
+            schema=data_schema,
+        )
+    else:
+        raise TypeError()
+
+
+# TODO: See comments about InputData above
+
+# _OutputDataType = TypeVar("_OutputDataType")
+# _OutputDataSchema = TypeVar("_OutputDataSchema")
+# _OutputDataLevel = TypeVar("_OutputDataLevel")
+
+
+# class OutputData(
+#     Generic[_OutputDataType, _OutputDataSchema, _OutputDataLevel],
+# ):
+#     def __init__(self):
+#         raise NotImplementedError()
+
+#     def __class_getitem__(cls, args) -> typing.GenericAlias:
+#         return Annotated[Iterable[pyarrow.RecordBatch], Output(*args)]
+
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class MethodPrototype(Generic[P, R]):
+    """A wrapper for Python functions that implement Methods that knows how to create an
+    appropriate MethodCreateRequest based on the function signature."""
+
+    def __init__(
+        self,
+        f: Callable[P, R],
+        *,
+        scope: MethodScope,
+        description: Optional[str] = None,
+    ):
+        self.f = f
+        self.scope = scope
+        self.description = description
+        # This is similar to doing @functools.wraps() but it works with
+        # function objects
+        functools.update_wrapper(self, f)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        return self.f(*args, **kwargs)
+
+    def create_request(
+        self, *, account: str, modules: list[str]
+    ) -> MethodCreateRequest:
+        """Create a MethodCreateRequest for the wrapped function.
+
+        :param account: The .account field for the request
+        :param modules: The .modules field for the request. This should include at least
+            the module that contains the wrapped function.
+        """
+        name = self.f.__name__
+        hints = typing.get_type_hints(self.f, include_extras=True)
+
+        parameters: list[MethodParameter] = []
+        inputs: list[MethodInput] = []
+        output: Optional[MethodOutput] = None
+        for k, v in hints.items():
+            annotation = None
+            if metadata := getattr(v, "__metadata__", None):
+                for m in metadata:
+                    if isinstance(m, DataAnnotation):
+                        annotation = m
+                        break
+            if k == "return":
+                if annotation is None:
+                    continue
+                if annotation.level is None:
+                    raise ValueError("Must specify .level for Output")
+                if annotation.schema is None:
+                    raise ValueError("Must specify .schema for Output")
+                output = MethodOutput(
+                    kind=MethodOutputKind(annotation.kind),
+                    measurement=MeasurementSpec(
+                        name=name,
+                        description=self.description,
+                        level=MeasurementLevel(annotation.level),
+                        schema=annotation.schema,
+                    ),
+                )
+            elif annotation is None:
+                parameters.append(MethodParameter(keyword=k))
+            else:
+                inputs.append(
+                    MethodInput(kind=MethodInputKind(annotation.kind), keyword=k)
+                )
+
+        if output is None:
+            raise TypeError("Return type must be annotated with Output()")
+
+        return MethodCreateRequest(
+            account=account,
+            modules=modules,
+            name=name,
+            scope=self.scope,
+            description=self.description,
+            implementation=MethodImplementation(
+                kind=MethodImplementationKind.PythonFunction,
+                pythonFunction=MethodImplementationPythonFunction(
+                    fullyQualifiedName=".".join(_fqn(self.f))
+                ),
+            ),
+            parameters=parameters,
+            inputs=inputs,
+            output=output,
+        )
+
+
+def method(
+    *, scope: MethodScope, description: Optional[str] = None
+) -> Callable[[Callable[P, R]], MethodPrototype[P, R]]:
+    """Use this decorator to indicate that a Python function implements a
+    Dyff Method. This should be used in conjunction with appropriate type
+    annotations, e.g.::
+
+        @method
+        def my_method(
+            arg: str,
+            data: Annotated[pyarrow.dataset.Dataset, Input(Evaluation)]
+        ) -> Annotated[
+            Iterable[pyarrow.RecordBatch],
+            Output(Measurement, schema=MyPydanticType, level=MeasurementLevel.Instance)
+        ]:
+            ...
+
+    The wrapped function will be an instance of MethodPrototype, and you can
+    use its .create_request() member function to create an appropriate
+    MethodCreateRequest for the wrapped function.
+
+    :param scope: The .scope field for the Method
+    :param description: The .description field for the Method. If not specified,
+        the docstring of the wrapped function will be used.
+    :return: A decorator that returns a MethodPrototype.
+    """
+
+    def decorator(f: Callable[P, R]) -> MethodPrototype[P, R]:
+        nonlocal description
+        if description is None:
+            description = f.__doc__
+        return MethodPrototype(f, scope=scope, description=description)
+
+    return decorator
+
+
+def method_request(
+    f: MethodPrototype, *, account: str, modules: list[str]
+) -> MethodCreateRequest:
+    return f.create_request(account=account, modules=modules)
+
+
+__all__ = [
+    "Input",
+    # "InputData",
+    "MethodPrototype",
+    "Output",
+    # "OutputData",
+    "method",
+    "method_request",
+]

dyff/schema/v0/r1/base.py

@@ -558,11 +558,11 @@ def list_(
         return pydantic.conlist(item_type, min_items=list_size, max_items=list_size)
 
 
-# mypy gets confused because 'dict' is the name of a method in DyffDefaultSerializers
+# mypy gets confused because 'dict' is the name of a method in DyffBaseModel
 _ModelAsDict = dict[str, Any]
 
 
-class DyffDefaultSerializers(pydantic.BaseModel):
+class DyffBaseModel(pydantic.BaseModel):
     """This must be the base class for *all pydantic models* in the Dyff schema.
 
     Overrides serialization functions to serialize by alias, so that "round-trip"
@@ -571,6 +571,9 @@ class DyffDefaultSerializers(pydantic.BaseModel):
     Python reserved words like 'bytes' as field names.
     """
 
+    class Config:
+        extra = pydantic.Extra.forbid
+
     def dict(self, *, by_alias: bool = True, **kwargs) -> _ModelAsDict:
         return super().dict(by_alias=by_alias, **kwargs)
 
@@ -603,10 +606,10 @@ class DyffDefaultSerializers(pydantic.BaseModel):
 # don't have timezones set currently for historical reasons. It's actually
 # better if all datetimes in the system are UTC, so that their JSON
 # representations (i.e., isoformat strings) are well-ordered.
-class DyffSchemaBaseModel(DyffDefaultSerializers):
+class DyffSchemaBaseModel(DyffBaseModel):
     """This should be the base class for *almost all* non-request models in the Dyff
     schema. Models that do not inherit from this class *must* still inherit from
-    DyffDefaultSerializers.
+    DyffBaseModel.
 
     Adds a root validator to ensure that all datetime fields are represented in the UTC
     timezone. This is necessary to avoid errors when comparing "naive" and "aware"
@@ -614,9 +617,6 @@ class DyffSchemaBaseModel(DyffDefaultSerializers):
     datetimes are well-ordered.
     """
 
-    class Config:
-        extra = pydantic.Extra.forbid
-
     @pydantic.root_validator
     def _ensure_datetime_timezone_utc(cls, values):
         update = {}
@@ -633,7 +633,7 @@ class DyffSchemaBaseModel(DyffDefaultSerializers):
 __all__ = [
     "DTYPE",
     "DType",
-    "DyffDefaultSerializers",
+    "DyffBaseModel",
     "DyffSchemaBaseModel",
     "FixedWidthFloat",
     "FixedWidthInt",

dyff/schema/v0/r1/platform.py

@@ -885,6 +885,7 @@ class Dataset(DyffEntity, DatasetBase):
 class ModelSourceKinds(str, enum.Enum):
     GitLFS = "GitLFS"
     HuggingFaceHub = "HuggingFaceHub"
+    Mock = "Mock"
     OpenLLM = "OpenLLM"
     Upload = "Upload"
 
@@ -967,12 +968,14 @@ class ModelResources(DyffSchemaBaseModel):
 
 
 class ModelStorageMedium(str, enum.Enum):
+    Mock = "Mock"
     ObjectStorage = "ObjectStorage"
     PersistentVolume = "PersistentVolume"
 
 
 class ModelArtifactKind(str, enum.Enum):
     HuggingFaceCache = "HuggingFaceCache"
+    Mock = "Mock"
 
 
 class ModelArtifactHuggingFaceCache(DyffSchemaBaseModel):
@@ -990,7 +993,7 @@ class ModelArtifact(DyffSchemaBaseModel):
         description="How the model data is represented"
     )
     huggingFaceCache: Optional[ModelArtifactHuggingFaceCache] = pydantic.Field(
-        description="Model stored in a HuggingFace cache"
+        default=None, description="Model stored in a HuggingFace cache"
     )
 
 

dyff/schema/v0/r1/requests.py

@@ -19,7 +19,7 @@ from typing import Optional, Union
 
 import pydantic
 
-from .base import DyffDefaultSerializers
+from .base import DyffBaseModel
 from .platform import (
     AnalysisBase,
     DatasetBase,
@@ -39,7 +39,7 @@ from .platform import (
 from .version import SchemaVersion
 
 
-class DyffRequestDefaultValidators(DyffDefaultSerializers):
+class DyffRequestDefaultValidators(DyffBaseModel):
     """This must be the base class for *all* request models in the Dyff schema.
 
     Adds a root validator to ensure that all user-provided datetime fields have a

{dyff_schema-0.10.3.dist-info → dyff_schema-0.12.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dyff-schema
-Version: 0.10.3
+Version: 0.12.0
 Summary: Data models for the Dyff AI auditing platform.
 Author-email: Digital Safety Research Institute <contact@dsri.org>
 License: Apache-2.0

{dyff_schema-0.10.3.dist-info → dyff_schema-0.12.0.dist-info}/RECORD

@@ -1,5 +1,6 @@
 dyff/schema/__init__.py,sha256=JcpxaRHNYgLjJWLjVayLlqacb2GX49Pazpwb8m-BctM,1031
 dyff/schema/adapters.py,sha256=YMTHv_2VlLGFp-Kqwa6H51hjffHmk8gXjZilHysIF5Q,123
+dyff/schema/annotations.py,sha256=nE6Jk1PLqlShj8uqjE_EzZC9zYnTDW5AVtQcjysiK8M,10018
 dyff/schema/base.py,sha256=jvaNtsSZyFfsdUZTcY_U-yfLY5_GyrMxSXhON2R9XR0,119
 dyff/schema/copydoc.py,sha256=B4ZRpQmbFxi-3l9LCHvaJiVKb9VxADgC5vey804Febc,1075
 dyff/schema/errors.py,sha256=ow3yiucU4wGkeLmapUURp3eyaebwGUwDaVTXpPcrA7M,1542
@@ -21,9 +22,9 @@ dyff/schema/io/vllm.py,sha256=2q05M_-lTzq9oywKXHPPpCFCSDVCSsRQqtmERzWTtio,123
 dyff/schema/v0/__init__.py,sha256=L5y8UhRnojerPYHumsxQJRcHCNz8Hj9NM8b47mewMNs,92
 dyff/schema/v0/r1/__init__.py,sha256=L5y8UhRnojerPYHumsxQJRcHCNz8Hj9NM8b47mewMNs,92
 dyff/schema/v0/r1/adapters.py,sha256=2t2oxsnGfSEDKKDIEYw4qqLXMH7qlFIwPVuLyUmbsHs,23552
-dyff/schema/v0/r1/base.py,sha256=l7RVpEvT8s3BZoaYGnHshbjNLRUIU2__JFE3R_hKsdY,19448
-dyff/schema/v0/r1/platform.py,sha256=QYsbEzo3o8dFqf85xw9f91EZ_6oU1J6N91L_nLAsjro,67231
-dyff/schema/v0/r1/requests.py,sha256=w1swbNCbmByQq5VEUpDtqrg7vd9UL_40unX3GhOb5ZM,10478
+dyff/schema/v0/r1/base.py,sha256=i7eOKXDGS8_J9k2aVObUTpSOnA8CAgRW7Quj1fSbyRg,19403
+dyff/schema/v0/r1/platform.py,sha256=lOjRZemR8tHRVlHl3UblSWVym6ff7s8bAyxnZ056pYU,67299
+dyff/schema/v0/r1/requests.py,sha256=3veChMw6QseXNvNBgqUGFq5fIK7S1JPOUMnVwstL5bI,10460
 dyff/schema/v0/r1/test.py,sha256=X6dUyVd5svcPCI-PBMOAqEfK9jv3bRDvkQTJzwS96c0,10720
 dyff/schema/v0/r1/version.py,sha256=isKAGuGxsdru8vDaYmI4YiZdJOu_wNxXK7u6QzD6FE4,392
 dyff/schema/v0/r1/dataset/__init__.py,sha256=LbVlkO2asyGYBKk2z49xjJYTM-pu9y9e4eQDXgTDLnM,2553
@@ -34,9 +35,9 @@ dyff/schema/v0/r1/dataset/text.py,sha256=nLIn91Zlt0tNdXUklSgjJ-kEDxoPX32ISLkiv2D
 dyff/schema/v0/r1/dataset/vision.py,sha256=aIe0fbfM_g3DsrDTdg2K803YKLjZBpurM_VJcJFuZLc,369
 dyff/schema/v0/r1/io/__init__.py,sha256=L5y8UhRnojerPYHumsxQJRcHCNz8Hj9NM8b47mewMNs,92
 dyff/schema/v0/r1/io/vllm.py,sha256=CUE9y8KthtUI7sD49S875rDmPvKotSXVIRaBS79aBZs,5320
-dyff_schema-0.10.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-dyff_schema-0.10.3.dist-info/METADATA,sha256=wic70dwTBnxwOHe7vC23RAeRrkESgu5-Qqqg-SgQUfk,3484
-dyff_schema-0.10.3.dist-info/NOTICE,sha256=YONACu0s_Ui6jNi-wtEsVQbTU1JIkh8wvLH6d1-Ni_w,43
-dyff_schema-0.10.3.dist-info/WHEEL,sha256=-oYQCr74JF3a37z2nRlQays_SX2MqOANoqVjBBAP2yE,91
-dyff_schema-0.10.3.dist-info/top_level.txt,sha256=9e3VVdeX73t_sUJOPQPCcGtYO1JhoErhHIi3WoWGcFI,5
-dyff_schema-0.10.3.dist-info/RECORD,,
+dyff_schema-0.12.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+dyff_schema-0.12.0.dist-info/METADATA,sha256=uK7or3ItWEusSeKVIO1ArFRPiopvSeuLqbPcc2CKhl8,3484
+dyff_schema-0.12.0.dist-info/NOTICE,sha256=YONACu0s_Ui6jNi-wtEsVQbTU1JIkh8wvLH6d1-Ni_w,43
+dyff_schema-0.12.0.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+dyff_schema-0.12.0.dist-info/top_level.txt,sha256=9e3VVdeX73t_sUJOPQPCcGtYO1JhoErhHIi3WoWGcFI,5
+dyff_schema-0.12.0.dist-info/RECORD,,

{dyff_schema-0.10.3.dist-info → dyff_schema-0.12.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (71.0.3)
+Generator: setuptools (71.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any