mlflow--mlflow/mlflow/entities/__init__.py

"""
The ``mlflow.entities`` module defines entities returned by the MLflow
`REST API <../rest-api.html>`_.
"""

from mlflow.entities.assessment import (
    Assessment,
    AssessmentError,
    AssessmentSource,
    AssessmentSourceType,
    Expectation,
    Feedback,
    IssueReference,
)
from mlflow.entities.dataset import Dataset
from mlflow.entities.dataset_input import DatasetInput
from mlflow.entities.dataset_record import DatasetRecord
from mlflow.entities.dataset_record_source import DatasetRecordSource, DatasetRecordSourceType
from mlflow.entities.dataset_summary import _DatasetSummary
from mlflow.entities.document import Document
from mlflow.entities.entity_type import EntityAssociationType
from mlflow.entities.experiment import Experiment
from mlflow.entities.experiment_tag import ExperimentTag
from mlflow.entities.file_info import FileInfo
from mlflow.entities.gateway_budget_policy import (
    BudgetAction,
    BudgetDuration,
    BudgetDurationUnit,
    BudgetTargetScope,
    BudgetUnit,
    GatewayBudgetPolicy,
)
from mlflow.entities.gateway_endpoint import (
    FallbackConfig,
    FallbackStrategy,
    GatewayEndpoint,
    GatewayEndpointBinding,
    GatewayEndpointModelConfig,
    GatewayEndpointModelMapping,
    GatewayEndpointTag,
    GatewayModelDefinition,
    GatewayModelLinkageType,
    GatewayResourceType,
    RoutingStrategy,
)
from mlflow.entities.gateway_guardrail import (
    GatewayGuardrail,
    GatewayGuardrailConfig,
    GuardrailAction,
    GuardrailStage,
)
from mlflow.entities.gateway_secrets import GatewaySecretInfo
from mlflow.entities.input_tag import InputTag
from mlflow.entities.issue import Issue, IssueSeverity, IssueStatus
from mlflow.entities.lifecycle_stage import LifecycleStage
from mlflow.entities.link import Link
from mlflow.entities.logged_model import LoggedModel
from mlflow.entities.logged_model_input import LoggedModelInput
from mlflow.entities.logged_model_output import LoggedModelOutput
from mlflow.entities.logged_model_parameter import LoggedModelParameter
from mlflow.entities.logged_model_status import LoggedModelStatus
from mlflow.entities.logged_model_tag import LoggedModelTag
from mlflow.entities.metric import Metric
from mlflow.entities.model_registry import Prompt
from mlflow.entities.param import Param
from mlflow.entities.run import Run
from mlflow.entities.run_data import RunData
from mlflow.entities.run_info import RunInfo
from mlflow.entities.run_inputs import RunInputs
from mlflow.entities.run_outputs import RunOutputs
from mlflow.entities.run_status import RunStatus
from mlflow.entities.run_tag import RunTag
from mlflow.entities.scorer import ScorerVersion
from mlflow.entities.session import Session
from mlflow.entities.source_type import SourceType
from mlflow.entities.span import LiveSpan, NoOpSpan, Span, SpanType
from mlflow.entities.span_event import SpanEvent
from mlflow.entities.span_log_level import SpanLogLevel
from mlflow.entities.span_status import SpanStatus, SpanStatusCode
from mlflow.entities.trace import Trace
from mlflow.entities.trace_data import TraceData
from mlflow.entities.trace_info import TraceInfo
from mlflow.entities.trace_location import (
    InferenceTableLocation,
    MlflowExperimentLocation,
    TraceLocation,
    TraceLocationType,
    UCSchemaLocation,
    UnityCatalog,
)
from mlflow.entities.trace_state import TraceState
from mlflow.entities.view_type import ViewType
from mlflow.entities.webhook import (
    Webhook,
    WebhookEvent,
    WebhookStatus,
    WebhookTestResult,
)
from mlflow.entities.workspace import TraceArchivalConfig, Workspace, WorkspaceDeletionMode

__all__ = [
    "Experiment",
    "ExperimentTag",
    "FileInfo",
    "Metric",
    "Param",
    "Prompt",
    "Run",
    "RunData",
    "RunInfo",
    "RunStatus",
    "RunTag",
    "ScorerVersion",
    "SourceType",
    "ViewType",
    "LifecycleStage",
    "Dataset",
    "InputTag",
    "Issue",
    "IssueSeverity",
    "IssueStatus",
    "DatasetInput",
    "RunInputs",
    "RunOutputs",
    "Link",
    "Span",
    "LiveSpan",
    "NoOpSpan",
    "SpanEvent",
    "SpanLogLevel",
    "SpanStatus",
    "SpanType",
    "Trace",
    "TraceData",
    "TraceInfo",
    "Session",
    "TraceLocation",
    "TraceLocationType",
    "MlflowExperimentLocation",
    "InferenceTableLocation",
    "UCSchemaLocation",
    "UnityCatalog",
    "TraceState",
    "SpanStatusCode",
    "_DatasetSummary",
    "LoggedModel",
    "LoggedModelInput",
    "LoggedModelOutput",
    "LoggedModelStatus",
    "LoggedModelTag",
    "LoggedModelParameter",
    "Document",
    "Assessment",
    "AssessmentError",
    "AssessmentSource",
    "AssessmentSourceType",
    "Expectation",
    "Feedback",
    "IssueReference",
    # Note: EvaluationDataset is intentionally excluded from __all__ to prevent
    # circular import issues during plugin registration. It can still be imported
    # explicitly via: from mlflow.entities import EvaluationDataset
    "DatasetRecord",
    "DatasetRecordSource",
    "DatasetRecordSourceType",
    "EntityAssociationType",
    "BudgetAction",
    "BudgetDuration",
    "BudgetDurationUnit",
    "BudgetTargetScope",
    "BudgetUnit",
    "FallbackConfig",
    "FallbackStrategy",
    "GatewayBudgetPolicy",
    "GatewayEndpoint",
    "GatewayEndpointBinding",
    "GatewayEndpointModelConfig",
    "GatewayEndpointModelMapping",
    "GatewayEndpointTag",
    "GatewayModelDefinition",
    "GatewayResourceType",
    "GatewaySecretInfo",
    "GatewayModelLinkageType",
    "RoutingStrategy",
    "Webhook",
    "WebhookEvent",
    "WebhookStatus",
    "WebhookTestResult",
    "TraceArchivalConfig",
    "Workspace",
    "WorkspaceDeletionMode",
    "GatewayGuardrail",
    "GatewayGuardrailConfig",
    "GuardrailAction",
    "GuardrailStage",
]


def __getattr__(name):
    """Lazy loading for EvaluationDataset to avoid circular imports."""
    if name == "EvaluationDataset":
        try:
            from mlflow.entities.evaluation_dataset import EvaluationDataset

            return EvaluationDataset
        except ImportError:
            # EvaluationDataset requires mlflow.data which may not be available
            # in minimal installations like mlflow-tracing
            raise AttributeError(
                "EvaluationDataset is not available. It requires the mlflow.data module "
                "which is not included in this installation."
            )
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")