genesis-flow 1.0.0__py3-none-any.whl

mlflow/entities/assessment_source.py

@@ -0,0 +1,141 @@
+import warnings
+from dataclasses import asdict, dataclass
+from typing import Any
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.exceptions import MlflowException
+from mlflow.protos.assessments_pb2 import AssessmentSource as ProtoAssessmentSource
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.annotations import experimental
+
+
+@experimental(version="2.21.0")
+@dataclass
+class AssessmentSource(_MlflowObject):
+    """
+    Source of an assessment (human, LLM as a judge with GPT-4, etc).
+
+    When recording an assessment, MLflow mandates providing a source information
+    to keep track of how the assessment is conducted.
+
+    Args:
+        source_type: The type of the assessment source. Must be one of the values in
+            the AssessmentSourceType enum.
+        source_id: An identifier for the source, e.g. user ID or LLM judge ID. If not
+            provided, the default value "default" is used.
+
+    Example:
+
+    Human annotation can be represented with a source type of "HUMAN":
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentSource, AssessmentSourceType
+
+        source = AssessmentSource(
+            source_type=AssessmentSourceType.HUMAN,
+            source_id="bob@example.com",
+        )
+
+    LLM-as-a-judge can be represented with a source type of "LLM_JUDGE":
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentSource, AssessmentSourceType
+
+        source = AssessmentSource(
+            source_type=AssessmentSourceType.LLM_JUDGE,
+            source_id="gpt-4o-mini",
+        )
+
+    Heuristic evaluation can be represented with a source type of "CODE":
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentSource, AssessmentSourceType
+
+        source = AssessmentSource(
+            source_type=AssessmentSourceType.CODE,
+            source_id="repo/evaluation_script.py",
+        )
+
+    To record more context about the assessment, you can use the `metadata` field of
+    the assessment logging APIs as well.
+    """
+
+    source_type: str
+    source_id: str = "default"
+
+    def __post_init__(self):
+        # Perform the standardization on source_type after initialization
+        self.source_type = AssessmentSourceType._standardize(self.source_type)
+
+    def to_dictionary(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dictionary(cls, source_dict: dict[str, Any]) -> "AssessmentSource":
+        return cls(**source_dict)
+
+    def to_proto(self):
+        source = ProtoAssessmentSource()
+        source.source_type = ProtoAssessmentSource.SourceType.Value(self.source_type)
+        if self.source_id is not None:
+            source.source_id = self.source_id
+        return source
+
+    @classmethod
+    def from_proto(cls, proto):
+        return AssessmentSource(
+            source_type=AssessmentSourceType.from_proto(proto.source_type),
+            source_id=proto.source_id if proto.source_id else None,
+        )
+
+
+@experimental(version="2.21.0")
+class AssessmentSourceType:
+    SOURCE_TYPE_UNSPECIFIED = "SOURCE_TYPE_UNSPECIFIED"
+    LLM_JUDGE = "LLM_JUDGE"
+    AI_JUDGE = "AI_JUDGE"  # Deprecated, use LLM_JUDGE instead
+    HUMAN = "HUMAN"
+    CODE = "CODE"
+    _SOURCE_TYPES = [SOURCE_TYPE_UNSPECIFIED, LLM_JUDGE, HUMAN, CODE]
+
+    def __init__(self, source_type: str):
+        self._source_type = AssessmentSourceType._parse(source_type)
+
+    @staticmethod
+    def _parse(source_type: str) -> str:
+        source_type = source_type.upper()
+
+        # Backwards compatibility shim for mlflow.evaluations.AssessmentSourceType
+        if source_type == AssessmentSourceType.AI_JUDGE:
+            warnings.warn(
+                "AI_JUDGE is deprecated. Use LLM_JUDGE instead.",
+                DeprecationWarning,
+            )
+            source_type = AssessmentSourceType.LLM_JUDGE
+
+        if source_type not in AssessmentSourceType._SOURCE_TYPES:
+            raise MlflowException(
+                message=(
+                    f"Invalid assessment source type: {source_type}. "
+                    f"Valid source types: {AssessmentSourceType._SOURCE_TYPES}"
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        return source_type
+
+    def __str__(self):
+        return self._source_type
+
+    @staticmethod
+    def _standardize(source_type: str) -> str:
+        return str(AssessmentSourceType(source_type))
+
+    @classmethod
+    def from_proto(cls, proto_source_type) -> str:
+        return ProtoAssessmentSource.SourceType.Name(proto_source_type)

mlflow/entities/dataset.py

@@ -0,0 +1,92 @@
+from typing import Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import Dataset as ProtoDataset
+
+
+class Dataset(_MlflowObject):
+    """Dataset object associated with an experiment."""
+
+    def __init__(
+        self,
+        name: str,
+        digest: str,
+        source_type: str,
+        source: str,
+        schema: Optional[str] = None,
+        profile: Optional[str] = None,
+    ) -> None:
+        self._name = name
+        self._digest = digest
+        self._source_type = source_type
+        self._source = source
+        self._schema = schema
+        self._profile = profile
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def name(self) -> str:
+        """String name of the dataset."""
+        return self._name
+
+    @property
+    def digest(self) -> str:
+        """String digest of the dataset."""
+        return self._digest
+
+    @property
+    def source_type(self) -> str:
+        """String source_type of the dataset."""
+        return self._source_type
+
+    @property
+    def source(self) -> str:
+        """String source of the dataset."""
+        return self._source
+
+    @property
+    def schema(self) -> str:
+        """String schema of the dataset."""
+        return self._schema
+
+    @property
+    def profile(self) -> str:
+        """String profile of the dataset."""
+        return self._profile
+
+    def to_proto(self):
+        dataset = ProtoDataset()
+        dataset.name = self.name
+        dataset.digest = self.digest
+        dataset.source_type = self.source_type
+        dataset.source = self.source
+        if self.schema:
+            dataset.schema = self.schema
+        if self.profile:
+            dataset.profile = self.profile
+        return dataset
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            proto.name,
+            proto.digest,
+            proto.source_type,
+            proto.source,
+            proto.schema if proto.HasField("schema") else None,
+            proto.profile if proto.HasField("profile") else None,
+        )
+
+    def to_dictionary(self):
+        return {
+            "name": self.name,
+            "digest": self.digest,
+            "source_type": self.source_type,
+            "source": self.source,
+            "schema": self.schema,
+            "profile": self.profile,
+        }

mlflow/entities/dataset_input.py

@@ -0,0 +1,51 @@
+from typing import Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.dataset import Dataset
+from mlflow.entities.input_tag import InputTag
+from mlflow.protos.service_pb2 import DatasetInput as ProtoDatasetInput
+
+
+class DatasetInput(_MlflowObject):
+    """DatasetInput object associated with an experiment."""
+
+    def __init__(self, dataset: Dataset, tags: Optional[list[InputTag]] = None) -> None:
+        self._dataset = dataset
+        self._tags = tags or []
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    def _add_tag(self, tag: InputTag) -> None:
+        self._tags.append(tag)
+
+    @property
+    def tags(self) -> list[InputTag]:
+        """Array of input tags."""
+        return self._tags
+
+    @property
+    def dataset(self) -> Dataset:
+        """Dataset."""
+        return self._dataset
+
+    def to_proto(self):
+        dataset_input = ProtoDatasetInput()
+        dataset_input.tags.extend([tag.to_proto() for tag in self.tags])
+        dataset_input.dataset.MergeFrom(self.dataset.to_proto())
+        return dataset_input
+
+    @classmethod
+    def from_proto(cls, proto):
+        dataset_input = cls(Dataset.from_proto(proto.dataset))
+        for input_tag in proto.tags:
+            dataset_input._add_tag(InputTag.from_proto(input_tag))
+        return dataset_input
+
+    def to_dictionary(self):
+        return {
+            "dataset": self.dataset.to_dictionary(),
+            "tags": {tag.key: tag.value for tag in self.tags},
+        }

mlflow/entities/dataset_summary.py

@@ -0,0 +1,62 @@
+from mlflow.protos.service_pb2 import DatasetSummary
+
+
+class _DatasetSummary:
+    """
+    DatasetSummary object.
+
+    This is used to return a list of dataset summaries across one or more experiments in the UI.
+    """
+
+    def __init__(self, experiment_id, name, digest, context):
+        self._experiment_id = experiment_id
+        self._name = name
+        self._digest = digest
+        self._context = context
+
+    def __eq__(self, other) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def experiment_id(self):
+        return self._experiment_id
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def digest(self):
+        return self._digest
+
+    @property
+    def context(self):
+        return self._context
+
+    def to_dict(self):
+        return {
+            "experiment_id": self.experiment_id,
+            "name": self.name,
+            "digest": self.digest,
+            "context": self.context,
+        }
+
+    def to_proto(self):
+        dataset_summary = DatasetSummary()
+        dataset_summary.experiment_id = self.experiment_id
+        dataset_summary.name = self.name
+        dataset_summary.digest = self.digest
+        if self.context:
+            dataset_summary.context = self.context
+        return dataset_summary
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            experiment_id=proto.experiment_id,
+            name=proto.name,
+            digest=proto.digest,
+            context=proto.context,
+        )

mlflow/entities/document.py

@@ -0,0 +1,48 @@
+from copy import deepcopy
+from dataclasses import asdict, dataclass, field
+from typing import Any, Optional
+
+
+@dataclass
+class Document:
+    """
+    An entity used in MLflow Tracing to represent retrieved documents in a RETRIEVER span.
+
+    Args:
+        page_content: The content of the document.
+        metadata: A dictionary of metadata associated with the document.
+        id: The ID of the document.
+    """
+
+    page_content: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+    id: Optional[str] = None
+
+    @classmethod
+    def from_langchain_document(cls, document):
+        # older versions of langchain do not have the id attribute
+        id = getattr(document, "id", None)
+
+        return cls(
+            page_content=document.page_content,
+            metadata=deepcopy(document.metadata),
+            id=id,
+        )
+
+    @classmethod
+    def from_llama_index_node_with_score(cls, node_with_score):
+        metadata = {
+            "score": node_with_score.get_score(),
+            # update after setting score so that it can be
+            # overridden if the user wishes to do so
+            **deepcopy(node_with_score.metadata),
+        }
+
+        return cls(
+            page_content=node_with_score.get_content(),
+            metadata=metadata,
+            id=node_with_score.node_id,
+        )
+
+    def to_dict(self):
+        return asdict(self)

mlflow/entities/experiment.py

@@ -0,0 +1,109 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.experiment_tag import ExperimentTag
+from mlflow.protos.service_pb2 import Experiment as ProtoExperiment
+from mlflow.protos.service_pb2 import ExperimentTag as ProtoExperimentTag
+
+
+class Experiment(_MlflowObject):
+    """
+    Experiment object.
+    """
+
+    DEFAULT_EXPERIMENT_NAME = "Default"
+
+    def __init__(
+        self,
+        experiment_id,
+        name,
+        artifact_location,
+        lifecycle_stage,
+        tags=None,
+        creation_time=None,
+        last_update_time=None,
+    ):
+        super().__init__()
+        self._experiment_id = experiment_id
+        self._name = name
+        self._artifact_location = artifact_location
+        self._lifecycle_stage = lifecycle_stage
+        self._tags = {tag.key: tag.value for tag in (tags or [])}
+        self._creation_time = creation_time
+        self._last_update_time = last_update_time
+
+    @property
+    def experiment_id(self):
+        """String ID of the experiment."""
+        return self._experiment_id
+
+    @property
+    def name(self):
+        """String name of the experiment."""
+        return self._name
+
+    def _set_name(self, new_name):
+        self._name = new_name
+
+    @property
+    def artifact_location(self):
+        """String corresponding to the root artifact URI for the experiment."""
+        return self._artifact_location
+
+    @property
+    def lifecycle_stage(self):
+        """Lifecycle stage of the experiment. Can either be 'active' or 'deleted'."""
+        return self._lifecycle_stage
+
+    @property
+    def tags(self):
+        """Tags that have been set on the experiment."""
+        return self._tags
+
+    def _add_tag(self, tag):
+        self._tags[tag.key] = tag.value
+
+    @property
+    def creation_time(self):
+        return self._creation_time
+
+    def _set_creation_time(self, creation_time):
+        self._creation_time = creation_time
+
+    @property
+    def last_update_time(self):
+        return self._last_update_time
+
+    def _set_last_update_time(self, last_update_time):
+        self._last_update_time = last_update_time
+
+    @classmethod
+    def from_proto(cls, proto):
+        experiment = cls(
+            proto.experiment_id,
+            proto.name,
+            proto.artifact_location,
+            proto.lifecycle_stage,
+            # `creation_time` and `last_update_time` were added in MLflow 1.29.0. Experiments
+            # created before this version don't have these fields and `proto.creation_time` and
+            # `proto.last_update_time` default to 0. We should only set `creation_time` and
+            # `last_update_time` if they are non-zero.
+            creation_time=proto.creation_time or None,
+            last_update_time=proto.last_update_time or None,
+        )
+        for proto_tag in proto.tags:
+            experiment._add_tag(ExperimentTag.from_proto(proto_tag))
+        return experiment
+
+    def to_proto(self):
+        experiment = ProtoExperiment()
+        experiment.experiment_id = self.experiment_id
+        experiment.name = self.name
+        experiment.artifact_location = self.artifact_location
+        experiment.lifecycle_stage = self.lifecycle_stage
+        if self.creation_time:
+            experiment.creation_time = self.creation_time
+        if self.last_update_time:
+            experiment.last_update_time = self.last_update_time
+        experiment.tags.extend(
+            [ProtoExperimentTag(key=key, value=val) for key, val in self._tags.items()]
+        )
+        return experiment

mlflow/entities/experiment_tag.py

@@ -0,0 +1,35 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import ExperimentTag as ProtoExperimentTag
+
+
+class ExperimentTag(_MlflowObject):
+    """Tag object associated with an experiment."""
+
+    def __init__(self, key, value):
+        self._key = key
+        self._value = value
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def key(self):
+        """String name of the tag."""
+        return self._key
+
+    @property
+    def value(self):
+        """String value of the tag."""
+        return self._value
+
+    def to_proto(self):
+        param = ProtoExperimentTag()
+        param.key = self.key
+        param.value = self.value
+        return param
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)

mlflow/entities/file_info.py

@@ -0,0 +1,45 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import FileInfo as ProtoFileInfo
+
+
+class FileInfo(_MlflowObject):
+    """
+    Metadata about a file or directory.
+    """
+
+    def __init__(self, path, is_dir, file_size):
+        self._path = path
+        self._is_dir = is_dir
+        self._bytes = file_size
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def path(self):
+        """String path of the file or directory."""
+        return self._path
+
+    @property
+    def is_dir(self):
+        """Whether the FileInfo corresponds to a directory."""
+        return self._is_dir
+
+    @property
+    def file_size(self):
+        """Size of the file or directory. If the FileInfo is a directory, returns None."""
+        return self._bytes
+
+    def to_proto(self):
+        proto = ProtoFileInfo()
+        proto.path = self.path
+        proto.is_dir = self.is_dir
+        if self.file_size:
+            proto.file_size = self.file_size
+        return proto
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.path, proto.is_dir, proto.file_size)

mlflow/entities/input_tag.py

@@ -0,0 +1,35 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import InputTag as ProtoInputTag
+
+
+class InputTag(_MlflowObject):
+    """Input tag object associated with a dataset."""
+
+    def __init__(self, key: str, value: str) -> None:
+        self._key = key
+        self._value = value
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def key(self) -> str:
+        """String name of the input tag."""
+        return self._key
+
+    @property
+    def value(self) -> str:
+        """String value of the input tag."""
+        return self._value
+
+    def to_proto(self):
+        tag = ProtoInputTag()
+        tag.key = self.key
+        tag.value = self.value
+        return tag
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)

mlflow/entities/lifecycle_stage.py

@@ -0,0 +1,35 @@
+from mlflow.entities.view_type import ViewType
+from mlflow.exceptions import MlflowException
+
+
+class LifecycleStage:
+    ACTIVE = "active"
+    DELETED = "deleted"
+    _VALID_STAGES = {ACTIVE, DELETED}
+
+    @classmethod
+    def view_type_to_stages(cls, view_type=ViewType.ALL):
+        stages = []
+        if view_type == ViewType.ACTIVE_ONLY or view_type == ViewType.ALL:
+            stages.append(cls.ACTIVE)
+        if view_type == ViewType.DELETED_ONLY or view_type == ViewType.ALL:
+            stages.append(cls.DELETED)
+        return stages
+
+    @classmethod
+    def is_valid(cls, lifecycle_stage):
+        return lifecycle_stage in cls._VALID_STAGES
+
+    @classmethod
+    def matches_view_type(cls, view_type, lifecycle_stage):
+        if not cls.is_valid(lifecycle_stage):
+            raise MlflowException(f"Invalid lifecycle stage '{lifecycle_stage}'")
+
+        if view_type == ViewType.ALL:
+            return True
+        elif view_type == ViewType.ACTIVE_ONLY:
+            return lifecycle_stage == LifecycleStage.ACTIVE
+        elif view_type == ViewType.DELETED_ONLY:
+            return lifecycle_stage == LifecycleStage.DELETED
+        else:
+            raise MlflowException(f"Invalid view type '{view_type}'")