star

venv/lib/python3.9/site-packages/mlflow/entities/__init__.py

@@ -0,0 +1,74 @@
+"""
+The ``mlflow.entities`` module defines entities returned by the MLflow
+`REST API <../rest-api.html>`_.
+"""
+
+from mlflow.entities.assessment import (
+    Assessment,
+    AssessmentError,
+    AssessmentSource,
+    AssessmentSourceType,
+)
+from mlflow.entities.dataset import Dataset
+from mlflow.entities.dataset_input import DatasetInput
+from mlflow.entities.dataset_summary import _DatasetSummary
+from mlflow.entities.document import Document
+from mlflow.entities.experiment import Experiment
+from mlflow.entities.experiment_tag import ExperimentTag
+from mlflow.entities.file_info import FileInfo
+from mlflow.entities.input_tag import InputTag
+from mlflow.entities.lifecycle_stage import LifecycleStage
+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_status import RunStatus
+from mlflow.entities.run_tag import RunTag
+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_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.view_type import ViewType
+
+__all__ = [
+    "Experiment",
+    "FileInfo",
+    "Metric",
+    "Param",
+    "Prompt",
+    "Run",
+    "RunData",
+    "RunInfo",
+    "RunStatus",
+    "RunTag",
+    "ExperimentTag",
+    "SourceType",
+    "ViewType",
+    "LifecycleStage",
+    "Dataset",
+    "InputTag",
+    "DatasetInput",
+    "RunInputs",
+    "Span",
+    "LiveSpan",
+    "NoOpSpan",
+    "SpanEvent",
+    "SpanStatus",
+    "SpanType",
+    "Trace",
+    "TraceData",
+    "TraceInfo",
+    "SpanStatusCode",
+    "_DatasetSummary",
+    "Document",
+    "Assessment",
+    "AssessmentError",
+    "AssessmentSource",
+    "AssessmentSourceType",
+]

venv/lib/python3.9/site-packages/mlflow/entities/_mlflow_object.py

@@ -0,0 +1,52 @@
+import pprint
+from abc import abstractmethod
+
+
+class _MlflowObject:
+    def __iter__(self):
+        # Iterate through list of properties and yield as key -> value
+        for prop in self._properties():
+            yield prop, self.__getattribute__(prop)
+
+    @classmethod
+    def _get_properties_helper(cls):
+        return sorted([p for p in cls.__dict__ if isinstance(getattr(cls, p), property)])
+
+    @classmethod
+    def _properties(cls):
+        return cls._get_properties_helper()
+
+    @classmethod
+    @abstractmethod
+    def from_proto(cls, proto):
+        pass
+
+    @classmethod
+    def from_dictionary(cls, the_dict):
+        filtered_dict = {key: value for key, value in the_dict.items() if key in cls._properties()}
+        return cls(**filtered_dict)
+
+    def __repr__(self):
+        return to_string(self)
+
+
+def to_string(obj):
+    return _MlflowObjectPrinter().to_string(obj)
+
+
+def get_classname(obj):
+    return type(obj).__name__
+
+
+class _MlflowObjectPrinter:
+    def __init__(self):
+        super().__init__()
+        self.printer = pprint.PrettyPrinter()
+
+    def to_string(self, obj):
+        if isinstance(obj, _MlflowObject):
+            return f"<{get_classname(obj)}: {self._entity_to_string(obj)}>"
+        return self.printer.pformat(obj)
+
+    def _entity_to_string(self, entity):
+        return ", ".join([f"{key}={self.to_string(value)}" for key, value in entity])

venv/lib/python3.9/site-packages/mlflow/entities/assessment.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Optional, Union
+
+from google.protobuf.json_format import MessageToDict, ParseDict
+from google.protobuf.struct_pb2 import Value
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.assessment_error import AssessmentError
+from mlflow.entities.assessment_source import AssessmentSource, AssessmentSourceType  # noqa: F401
+from mlflow.exceptions import MlflowException
+from mlflow.protos.assessments_pb2 import Assessment as ProtoAssessment
+from mlflow.protos.assessments_pb2 import Expectation as ProtoExpectation
+from mlflow.protos.assessments_pb2 import Feedback as ProtoFeedback
+from mlflow.utils.annotations import experimental
+from mlflow.utils.proto_json_utils import parse_pb_value, set_pb_value
+
+# Assessment value should be one of the following types:
+# - float
+# - int
+# - str
+# - bool
+# - list of values of the same types as above
+# - dict with string keys and values of the same types as above
+PbValueType = Union[float, int, str, bool]
+AssessmentValueType = Union[PbValueType, dict[str, PbValueType], list[PbValueType]]
+
+
+@experimental
+@dataclass
+class Assessment(_MlflowObject):
+    """
+    An abstraction for annotating a trace. An Assessment should be one of the following types:
+
+    - Expectations: A label that represents the expected value for a particular operation.
+        For example, an expected answer for a user question from a chatbot.
+    - Feedback: A label that represents the feedback on the quality of the operation.
+        Feedback can come from different sources, such as human judges, heuristic scorers,
+        or LLM-as-a-Judge.
+
+    You can log an assessment to a trace using the :py:func:`mlflow.log_expectation` or
+    :py:func:`mlflow.log_feedback` functions.
+
+    Args:
+        name: The name of the assessment.
+        source: The source of the assessment.
+        trace_id: The ID of the trace associated with the assessment. If unset, the assessment
+            is not associated with any trace yet.
+        expectation: The expectation value of the assessment.
+        feedback: The feedback value of the assessment.  Only one of `expectation` or `feedback`
+            should be specified.
+        rationale: The rationale / justification for the assessment.
+        metadata: The metadata associated with the assessment.
+        span_id: The ID of the span associated with the assessment, if the assessment should
+            be associated with a particular span in the trace.
+        create_time_ms: The creation time of the assessment in milliseconds. If unset, the
+            current time is used.
+        last_update_time_ms: The last update time of the assessment in milliseconds.
+            If unset, the current time is used.
+        assessment_id: The ID of the assessment. This must be generated in the backend.
+    """
+
+    name: str
+    source: AssessmentSource
+    # NB: The trace ID is optional because the assessment object itself may be created
+    #   standalone. For example, a custom metric function returns an assessment object
+    #   without a trace ID. That said, the trace ID is required when logging the
+    #   assessment to a trace in the backend eventually.
+    #   https://docs.databricks.com/aws/en/generative-ai/agent-evaluation/custom-metrics#-metric-decorator
+    trace_id: Optional[str] = None
+    expectation: Optional[Expectation] = None
+    feedback: Optional[Feedback] = None
+    rationale: Optional[str] = None
+    metadata: Optional[dict[str, str]] = None
+    span_id: Optional[str] = None
+    create_time_ms: Optional[int] = None
+    last_update_time_ms: Optional[int] = None
+    # NB: The assessment ID should always be generated in the backend. The CreateAssessment
+    #   backend API asks for an incomplete Assessment object without an ID and returns a
+    #   complete one with assessment_id, so the ID is Optional in the constructor here.
+    assessment_id: Optional[str] = None
+    # Deprecated, use `error` in Feedback instead. Just kept for backward compatibility
+    # and will be removed in the 3.0.0 release.
+    error: Optional[AssessmentError] = None
+
+    def __post_init__(self):
+        if (self.expectation is not None) + (self.feedback is not None) != 1:
+            raise MlflowException.invalid_parameter_value(
+                "Exactly one of `expectation` or `feedback` should be specified.",
+            )
+
+        # Populate the error field to the feedback object
+        if self.error is not None:
+            if self.expectation is not None:
+                raise MlflowException.invalid_parameter_value(
+                    "Cannot set `error` when `expectation` is specified.",
+                )
+            if self.feedback is None:
+                raise MlflowException.invalid_parameter_value(
+                    "Cannot set `error` when `feedback` is not specified.",
+                )
+            self.feedback.error = self.error
+
+        # Set timestamp if not provided
+        current_time = int(time.time() * 1000)  # milliseconds
+        if self.create_time_ms is None:
+            self.create_time_ms = current_time
+        if self.last_update_time_ms is None:
+            self.last_update_time_ms = current_time
+
+    def to_proto(self):
+        assessment = ProtoAssessment()
+        assessment.assessment_name = self.name
+        assessment.trace_id = self.trace_id
+
+        assessment.source.CopyFrom(self.source.to_proto())
+
+        # Convert time in milliseconds to protobuf Timestamp
+        assessment.create_time.FromMilliseconds(self.create_time_ms)
+        assessment.last_update_time.FromMilliseconds(self.last_update_time_ms)
+
+        if self.span_id is not None:
+            assessment.span_id = self.span_id
+        if self.rationale is not None:
+            assessment.rationale = self.rationale
+        if self.assessment_id is not None:
+            assessment.assessment_id = self.assessment_id
+
+        if self.expectation is not None:
+            set_pb_value(assessment.expectation.value, self.expectation.value)
+        elif self.feedback is not None:
+            assessment.feedback.CopyFrom(self.feedback.to_proto())
+
+        if self.metadata:
+            assessment.metadata.update(self.metadata)
+
+        return assessment
+
+    @classmethod
+    def from_proto(cls, proto):
+        if proto.WhichOneof("value") == "expectation":
+            expectation = Expectation(parse_pb_value(proto.expectation.value))
+            feedback = None
+        elif proto.WhichOneof("value") == "feedback":
+            expectation = None
+            feedback = Feedback.from_proto(proto.feedback)
+        else:
+            expectation = None
+            feedback = None
+
+        # Convert ScalarMapContainer to a normal Python dict
+        metadata = dict(proto.metadata) if proto.metadata else None
+
+        return cls(
+            assessment_id=proto.assessment_id or None,
+            trace_id=proto.trace_id,
+            name=proto.assessment_name,
+            source=AssessmentSource.from_proto(proto.source),
+            create_time_ms=proto.create_time.ToMilliseconds(),
+            last_update_time_ms=proto.last_update_time.ToMilliseconds(),
+            expectation=expectation,
+            feedback=feedback,
+            rationale=proto.rationale or None,
+            metadata=metadata,
+            span_id=proto.span_id or None,
+        )
+
+    def to_dictionary(self):
+        return {
+            "assessment_id": self.assessment_id,
+            "trace_id": self.trace_id,
+            "name": self.name,
+            "source": self.source.to_dictionary(),
+            "create_time_ms": self.create_time_ms,
+            "last_update_time_ms": self.last_update_time_ms,
+            "expectation": self.expectation.to_dictionary() if self.expectation else None,
+            "feedback": self.feedback.to_dictionary() if self.feedback else None,
+            "rationale": self.rationale,
+            "metadata": self.metadata,
+            "span_id": self.span_id,
+        }
+
+
+@experimental
+@dataclass
+class Expectation(_MlflowObject):
+    """
+    Represents an expectation about the output of an operation, such as the expected response
+    that a generative AI application should provide to a particular user query.
+    """
+
+    value: AssessmentValueType
+
+    def to_proto(self):
+        expectation = ProtoExpectation()
+        expectation.value = self.value
+        return expectation
+
+    def to_dictionary(self):
+        return {"value": self.value}
+
+
+@experimental
+@dataclass
+class Feedback(_MlflowObject):
+    """
+    Represents feedback about the output of an operation. For example, if the response from a
+    generative AI application to a particular user query is correct, then a human or LLM judge
+    may provide feedback with the value ``"correct"``.
+    """
+
+    value: AssessmentValueType
+    error: Optional[AssessmentError] = None
+
+    def to_proto(self):
+        return ProtoFeedback(
+            value=ParseDict(self.value, Value(), ignore_unknown_fields=True),
+            error=self.error.to_proto() if self.error else None,
+        )
+
+    @classmethod
+    def from_proto(self, proto):
+        return Feedback(
+            value=MessageToDict(proto.value),
+            error=AssessmentError.from_proto(proto.error) if proto.HasField("error") else None,
+        )
+
+    def to_dictionary(self):
+        d = {"value": self.value}
+        if self.error:
+            d["error"] = self.error.to_dictionary()
+        return d

venv/lib/python3.9/site-packages/mlflow/entities/assessment_error.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.assessments_pb2 import AssessmentError as ProtoAssessmentError
+from mlflow.utils.annotations import experimental
+
+
+@experimental
+@dataclass
+class AssessmentError(_MlflowObject):
+    """
+    Error object representing any issues during generating the assessment.
+
+    For example, if the LLM-as-a-Judge fails to generate an feedback, you can
+    log an error with the error code and message as shown below:
+
+    .. code-block:: python
+
+        from mlflow.entities import AssessmentError
+
+        error = AssessmentError(
+            error_code="RATE_LIMIT_EXCEEDED",
+            error_message="Rate limit for the judge exceeded.",
+        )
+
+        mlflow.log_feedback(
+            trace_id="1234",
+            name="faithfulness",
+            source=AssessmentSourceType.LLM_JUDGE,
+            error=error,
+            # Skip setting value when an error is present
+        )
+
+    Args:
+        error_code: The error code.
+        error_message: The detailed error message. Optional.
+    """
+
+    error_code: str
+    error_message: Optional[str] = None
+
+    def to_proto(self):
+        error = ProtoAssessmentError()
+        error.error_code = self.error_code
+        if self.error_message:
+            error.error_message = self.error_message
+        return error
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            error_code=proto.error_code,
+            error_message=proto.error_message or None,
+        )
+
+    def to_dictionary(self):
+        return {"error_code": self.error_code, "error_message": self.error_message}
+
+    @classmethod
+    def from_dictionary(cls, error_dict):
+        return cls(**error_dict)

venv/lib/python3.9/site-packages/mlflow/entities/assessment_source.py

@@ -0,0 +1,96 @@
+import warnings
+from dataclasses import asdict, dataclass
+from typing import Any, Optional
+
+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
+@dataclass
+class AssessmentSource(_MlflowObject):
+    """
+    Source of an assessment (human, LLM as a judge with GPT-4, etc).
+
+    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.
+    """
+
+    source_type: str
+    source_id: Optional[str] = None
+
+    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
+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)

venv/lib/python3.9/site-packages/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,
+        }

venv/lib/python3.9/site-packages/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},
+        }

venv/lib/python3.9/site-packages/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,
+        )

venv/lib/python3.9/site-packages/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)

venv/lib/python3.9/site-packages/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

venv/lib/python3.9/site-packages/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)

venv/lib/python3.9/site-packages/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)

venv/lib/python3.9/site-packages/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)

venv/lib/python3.9/site-packages/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}'")

venv/lib/python3.9/site-packages/mlflow/entities/metric.py

@@ -0,0 +1,126 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.service_pb2 import Metric as ProtoMetric
+from mlflow.protos.service_pb2 import MetricWithRunId as ProtoMetricWithRunId
+
+
+class Metric(_MlflowObject):
+    """
+    Metric object.
+    """
+
+    def __init__(self, key, value, timestamp, step):
+        self._key = key
+        self._value = value
+        self._timestamp = timestamp
+        self._step = step
+
+    @property
+    def key(self):
+        """String key corresponding to the metric name."""
+        return self._key
+
+    @property
+    def value(self):
+        """Float value of the metric."""
+        return self._value
+
+    @property
+    def timestamp(self):
+        """Metric timestamp as an integer (milliseconds since the Unix epoch)."""
+        return self._timestamp
+
+    @property
+    def step(self):
+        """Integer metric step (x-coordinate)."""
+        return self._step
+
+    def to_proto(self):
+        metric = ProtoMetric()
+        metric.key = self.key
+        metric.value = self.value
+        metric.timestamp = self.timestamp
+        metric.step = self.step
+        return metric
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value, proto.timestamp, proto.step)
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self.__dict__ == __o.__dict__
+
+        return False
+
+    def __hash__(self):
+        return hash((self._key, self._value, self._timestamp, self._step))
+
+    def to_dictionary(self):
+        """
+        Convert the Metric object to a dictionary.
+
+        Returns:
+            dict: The Metric object represented as a dictionary.
+        """
+        return {
+            "key": self.key,
+            "value": self.value,
+            "timestamp": self.timestamp,
+            "step": self.step,
+        }
+
+    @classmethod
+    def from_dictionary(cls, metric_dict):
+        """
+        Create a Metric object from a dictionary.
+
+        Args:
+            metric_dict (dict): Dictionary containing metric information.
+
+        Returns:
+            Metric: The Metric object created from the dictionary.
+        """
+        required_keys = ["key", "value", "timestamp", "step"]
+        missing_keys = [key for key in required_keys if key not in metric_dict]
+        if missing_keys:
+            raise MlflowException(
+                f"Missing required keys {missing_keys} in metric dictionary",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        return cls(**metric_dict)
+
+
+class MetricWithRunId(Metric):
+    def __init__(self, metric: Metric, run_id):
+        super().__init__(
+            key=metric.key,
+            value=metric.value,
+            timestamp=metric.timestamp,
+            step=metric.step,
+        )
+        self._run_id = run_id
+
+    @property
+    def run_id(self):
+        return self._run_id
+
+    def to_dict(self):
+        return {
+            "key": self.key,
+            "value": self.value,
+            "timestamp": self.timestamp,
+            "step": self.step,
+            "run_id": self.run_id,
+        }
+
+    def to_proto(self):
+        metric = ProtoMetricWithRunId()
+        metric.key = self.key
+        metric.value = self.value
+        metric.timestamp = self.timestamp
+        metric.step = self.step
+        metric.run_id = self.run_id
+        return metric

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/__init__.py

@@ -0,0 +1,19 @@
+from mlflow.entities.model_registry.model_version import ModelVersion
+from mlflow.entities.model_registry.model_version_search import ModelVersionSearch
+from mlflow.entities.model_registry.model_version_tag import ModelVersionTag
+from mlflow.entities.model_registry.prompt import Prompt
+from mlflow.entities.model_registry.registered_model import RegisteredModel
+from mlflow.entities.model_registry.registered_model_alias import RegisteredModelAlias
+from mlflow.entities.model_registry.registered_model_search import RegisteredModelSearch
+from mlflow.entities.model_registry.registered_model_tag import RegisteredModelTag
+
+__all__ = [
+    "Prompt",
+    "RegisteredModel",
+    "ModelVersion",
+    "RegisteredModelAlias",
+    "RegisteredModelTag",
+    "ModelVersionTag",
+    "RegisteredModelSearch",
+    "ModelVersionSearch",
+]

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/_model_registry_entity.py

@@ -0,0 +1,13 @@
+from abc import abstractmethod
+
+from mlflow.entities._mlflow_object import _MlflowObject
+
+
+class _ModelRegistryEntity(_MlflowObject):
+    @classmethod
+    @abstractmethod
+    def from_proto(cls, proto):
+        pass
+
+    def __eq__(self, other):
+        return dict(self) == dict(other)

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version.py

@@ -0,0 +1,199 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.entities.model_registry.model_version_status import ModelVersionStatus
+from mlflow.entities.model_registry.model_version_tag import ModelVersionTag
+from mlflow.protos.model_registry_pb2 import ModelVersion as ProtoModelVersion
+from mlflow.protos.model_registry_pb2 import ModelVersionTag as ProtoModelVersionTag
+
+
+class ModelVersion(_ModelRegistryEntity):
+    """
+    MLflow entity for Model Version.
+    """
+
+    def __init__(
+        self,
+        name,
+        version,
+        creation_timestamp,
+        last_updated_timestamp=None,
+        description=None,
+        user_id=None,
+        current_stage=None,
+        source=None,
+        run_id=None,
+        status=ModelVersionStatus.to_string(ModelVersionStatus.READY),
+        status_message=None,
+        tags=None,
+        run_link=None,
+        aliases=None,
+    ):
+        super().__init__()
+        self._name = name
+        self._version = version
+        self._creation_time = creation_timestamp
+        self._last_updated_timestamp = last_updated_timestamp
+        self._description = description
+        self._user_id = user_id
+        self._current_stage = current_stage
+        self._source = source
+        self._run_id = run_id
+        self._run_link = run_link
+        self._status = status
+        self._status_message = status_message
+        self._tags = {tag.key: tag.value for tag in (tags or [])}
+        self._aliases = aliases or []
+
+    @property
+    def name(self):
+        """String. Unique name within Model Registry."""
+        return self._name
+
+    @name.setter
+    def name(self, new_name):
+        self._name = new_name
+
+    @property
+    def version(self):
+        """version"""
+        return self._version
+
+    @property
+    def creation_timestamp(self):
+        """Integer. Model version creation timestamp (milliseconds since the Unix epoch)."""
+        return self._creation_time
+
+    @property
+    def last_updated_timestamp(self):
+        """Integer. Timestamp of last update for this model version (milliseconds since the Unix
+        epoch).
+        """
+        return self._last_updated_timestamp
+
+    @last_updated_timestamp.setter
+    def last_updated_timestamp(self, updated_timestamp):
+        self._last_updated_timestamp = updated_timestamp
+
+    @property
+    def description(self):
+        """String. Description"""
+        return self._description
+
+    @description.setter
+    def description(self, description):
+        self._description = description
+
+    @property
+    def user_id(self):
+        """String. User ID that created this model version."""
+        return self._user_id
+
+    @property
+    def current_stage(self):
+        """String. Current stage of this model version."""
+        return self._current_stage
+
+    @current_stage.setter
+    def current_stage(self, stage):
+        self._current_stage = stage
+
+    @property
+    def source(self):
+        """String. Source path for the model."""
+        return self._source
+
+    @property
+    def run_id(self):
+        """String. MLflow run ID that generated this model."""
+        return self._run_id
+
+    @property
+    def run_link(self):
+        """String. MLflow run link referring to the exact run that generated this model version."""
+        return self._run_link
+
+    @property
+    def status(self):
+        """String. Current Model Registry status for this model."""
+        return self._status
+
+    @property
+    def status_message(self):
+        """String. Descriptive message for error status conditions."""
+        return self._status_message
+
+    @property
+    def tags(self):
+        """Dictionary of tag key (string) -> tag value for the current model version."""
+        return self._tags
+
+    @property
+    def aliases(self):
+        """List of aliases (string) for the current model version."""
+        return self._aliases
+
+    @aliases.setter
+    def aliases(self, aliases):
+        self._aliases = aliases
+
+    @classmethod
+    def _properties(cls):
+        # aggregate with base class properties since cls.__dict__ does not do it automatically
+        return sorted(cls._get_properties_helper())
+
+    def _add_tag(self, tag):
+        self._tags[tag.key] = tag.value
+
+    # proto mappers
+    @classmethod
+    def from_proto(cls, proto):
+        # input: mlflow.protos.model_registry_pb2.ModelVersion
+        # returns: ModelVersion entity
+        model_version = cls(
+            proto.name,
+            proto.version,
+            proto.creation_timestamp,
+            proto.last_updated_timestamp,
+            proto.description if proto.HasField("description") else None,
+            proto.user_id,
+            proto.current_stage,
+            proto.source,
+            proto.run_id if proto.HasField("run_id") else None,
+            ModelVersionStatus.to_string(proto.status),
+            proto.status_message if proto.HasField("status_message") else None,
+            run_link=proto.run_link,
+            aliases=proto.aliases,
+        )
+        for tag in proto.tags:
+            model_version._add_tag(ModelVersionTag.from_proto(tag))
+        return model_version
+
+    def to_proto(self):
+        # input: ModelVersion entity
+        # returns mlflow.protos.model_registry_pb2.ModelVersion
+        model_version = ProtoModelVersion()
+        model_version.name = self.name
+        model_version.version = str(self.version)
+        model_version.creation_timestamp = self.creation_timestamp
+        if self.last_updated_timestamp is not None:
+            model_version.last_updated_timestamp = self.last_updated_timestamp
+        if self.description is not None:
+            model_version.description = self.description
+        if self.user_id is not None:
+            model_version.user_id = self.user_id
+        if self.current_stage is not None:
+            model_version.current_stage = self.current_stage
+        if self.source is not None:
+            model_version.source = str(self.source)
+        if self.run_id is not None:
+            model_version.run_id = str(self.run_id)
+        if self.run_link is not None:
+            model_version.run_link = str(self.run_link)
+        if self.status is not None:
+            model_version.status = ModelVersionStatus.from_string(self.status)
+        if self.status_message:
+            model_version.status_message = self.status_message
+        model_version.tags.extend(
+            [ProtoModelVersionTag(key=key, value=value) for key, value in self._tags.items()]
+        )
+        model_version.aliases.extend(self.aliases)
+        return model_version

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_search.py

@@ -0,0 +1,25 @@
+from mlflow.entities.model_registry import ModelVersion
+
+
+class ModelVersionSearch(ModelVersion):
+    def __init__(self, *args, **kwargs):
+        kwargs["tags"] = []
+        kwargs["aliases"] = []
+        super().__init__(*args, **kwargs)
+
+    def tags(self):
+        raise Exception(
+            "UC Model Versions gathered through search_model_versions do not have tags. "
+            "Please use get_model_version to obtain an individual version's tags."
+        )
+
+    def aliases(self):
+        raise Exception(
+            "UC Model Versions gathered through search_model_versions do not have aliases. "
+            "Please use get_model_version to obtain an individual version's aliases."
+        )
+
+    def __eq__(self, other):
+        if type(other) in {type(self), ModelVersion}:
+            return self.__dict__ == other.__dict__
+        return False

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_stages.py

@@ -0,0 +1,25 @@
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+STAGE_NONE = "None"
+STAGE_STAGING = "Staging"
+STAGE_PRODUCTION = "Production"
+STAGE_ARCHIVED = "Archived"
+
+STAGE_DELETED_INTERNAL = "Deleted_Internal"
+
+ALL_STAGES = [STAGE_NONE, STAGE_STAGING, STAGE_PRODUCTION, STAGE_ARCHIVED]
+DEFAULT_STAGES_FOR_GET_LATEST_VERSIONS = [STAGE_STAGING, STAGE_PRODUCTION]
+_CANONICAL_MAPPING = {stage.lower(): stage for stage in ALL_STAGES}
+
+
+def get_canonical_stage(stage):
+    key = stage.lower()
+    if key not in _CANONICAL_MAPPING:
+        raise MlflowException(
+            "Invalid Model Version stage: {}. Value must be one of {}.".format(
+                stage, ", ".join(ALL_STAGES)
+            ),
+            INVALID_PARAMETER_VALUE,
+        )
+    return _CANONICAL_MAPPING[key]

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_status.py

@@ -0,0 +1,35 @@
+from mlflow.protos.model_registry_pb2 import ModelVersionStatus as ProtoModelVersionStatus
+
+
+class ModelVersionStatus:
+    """Enum for status of an :py:class:`mlflow.entities.model_registry.ModelVersion`."""
+
+    PENDING_REGISTRATION = ProtoModelVersionStatus.Value("PENDING_REGISTRATION")
+    FAILED_REGISTRATION = ProtoModelVersionStatus.Value("FAILED_REGISTRATION")
+    READY = ProtoModelVersionStatus.Value("READY")
+    _STRING_TO_STATUS = {
+        k: ProtoModelVersionStatus.Value(k) for k in ProtoModelVersionStatus.keys()
+    }
+    _STATUS_TO_STRING = {value: key for key, value in _STRING_TO_STATUS.items()}
+
+    @staticmethod
+    def from_string(status_str):
+        if status_str not in ModelVersionStatus._STRING_TO_STATUS:
+            raise Exception(
+                f"Could not get model version status corresponding to string {status_str}. "
+                f"Valid status strings: {list(ModelVersionStatus._STRING_TO_STATUS.keys())}"
+            )
+        return ModelVersionStatus._STRING_TO_STATUS[status_str]
+
+    @staticmethod
+    def to_string(status):
+        if status not in ModelVersionStatus._STATUS_TO_STRING:
+            raise Exception(
+                f"Could not get string corresponding to model version status {status}. "
+                f"Valid statuses: {list(ModelVersionStatus._STATUS_TO_STRING.keys())}"
+            )
+        return ModelVersionStatus._STATUS_TO_STRING[status]
+
+    @staticmethod
+    def all_status():
+        return list(ModelVersionStatus._STATUS_TO_STRING.keys())

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_tag.py

@@ -0,0 +1,35 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.protos.model_registry_pb2 import ModelVersionTag as ProtoModelVersionTag
+
+
+class ModelVersionTag(_ModelRegistryEntity):
+    """Tag object associated with a model version."""
+
+    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
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)
+
+    def to_proto(self):
+        tag = ProtoModelVersionTag()
+        tag.key = self.key
+        tag.value = self.value
+        return tag

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/prompt.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+import re
+from typing import Optional, Union
+
+from mlflow.entities.model_registry.model_version import ModelVersion
+from mlflow.entities.model_registry.model_version_tag import ModelVersionTag
+from mlflow.exceptions import MlflowException
+from mlflow.prompt.constants import (
+    IS_PROMPT_TAG_KEY,
+    PROMPT_ASSOCIATED_RUN_IDS_TAG_KEY,
+    PROMPT_TEMPLATE_VARIABLE_PATTERN,
+    PROMPT_TEXT_DISPLAY_LIMIT,
+    PROMPT_TEXT_TAG_KEY,
+)
+
+# Alias type
+PromptVersionTag = ModelVersionTag
+
+
+def _is_reserved_tag(key: str) -> bool:
+    return key in {IS_PROMPT_TAG_KEY, PROMPT_TEXT_TAG_KEY, PROMPT_ASSOCIATED_RUN_IDS_TAG_KEY}
+
+
+# Prompt is implemented as a special type of ModelVersion. MLflow stores both prompts
+# and model versions in the model registry as ModelVersion DB records, but distinguishes
+# them using the special tag "mlflow.prompt.is_prompt".
+class Prompt(ModelVersion):
+    """
+    An entity representing a prompt (template) for GenAI applications.
+
+    Args:
+        name: The name of the prompt.
+        version: The version number of the prompt.
+        template: The template text of the prompt. It can contain variables enclosed in
+            double curly braces, e.g. {{variable}}, which will be replaced with actual values
+            by the `format` method. MLflow use the same variable naming rules same as Jinja2
+            https://jinja.palletsprojects.com/en/stable/api/#notes-on-identifiers
+        commit_message: The commit message for the prompt version. Optional.
+        creation_timestamp: Timestamp of the prompt creation. Optional.
+        version_metadata: A dictionary of metadata associated with the **prompt version**.
+            This is useful for storing version-specific information, such as the author of
+            the changes. Optional.
+        prompt_tags: A dictionary of tags associated with the entire prompt. This is different
+            from the `version_metadata` as it is not tied to a specific version of the prompt.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        version: int,
+        template: str,
+        commit_message: Optional[str] = None,
+        creation_timestamp: Optional[int] = None,
+        version_metadata: Optional[dict[str, str]] = None,
+        prompt_tags: Optional[dict[str, str]] = None,
+        aliases: Optional[list[str]] = None,
+    ):
+        # Store template text as a tag
+        version_metadata = version_metadata or {}
+        version_metadata[PROMPT_TEXT_TAG_KEY] = template
+        version_metadata[IS_PROMPT_TAG_KEY] = "true"
+
+        super().__init__(
+            name=name,
+            version=version,
+            creation_timestamp=creation_timestamp,
+            description=commit_message,
+            # "version_metadata" is represented as ModelVersion tags.
+            tags=[ModelVersionTag(key=key, value=value) for key, value in version_metadata.items()],
+            aliases=aliases,
+        )
+
+        self._variables = set(PROMPT_TEMPLATE_VARIABLE_PATTERN.findall(self.template))
+
+        # Store the prompt-level tags (from RegisteredModel).
+        self._prompt_tags = prompt_tags or {}
+
+    def __repr__(self) -> str:
+        text = (
+            self.template[:PROMPT_TEXT_DISPLAY_LIMIT] + "..."
+            if len(self.template) > PROMPT_TEXT_DISPLAY_LIMIT
+            else self.template
+        )
+        return f"Prompt(name={self.name}, version={self.version}, template={text})"
+
+    @property
+    def template(self) -> str:
+        """
+        Return the template text of the prompt.
+        """
+        return self._tags[PROMPT_TEXT_TAG_KEY]
+
+    def to_single_brace_format(self) -> str:
+        """
+        Convert the template text to single brace format. This is useful for integrating with other
+        systems that use single curly braces for variable replacement, such as LangChain's prompt
+        template. Default is False.
+        """
+        t = self.template
+        for var in self.variables:
+            t = re.sub(r"\{\{\s*" + var + r"\s*\}\}", "{" + var + "}", t)
+        return t
+
+    @property
+    def variables(self) -> set[str]:
+        """
+        Return a list of variables in the template text.
+        The value must be enclosed in double curly braces, e.g. {{variable}}.
+        """
+        return self._variables
+
+    @property
+    def commit_message(self) -> Optional[str]:
+        """
+        Return the commit message of the prompt version.
+        """
+        return self.description  # inherited from ModelVersion
+
+    @property
+    def version_metadata(self) -> dict[str, str]:
+        """Return the tags of the prompt as a dictionary."""
+        # Remove the prompt text tag as it should not be user-facing
+        return {key: value for key, value in self._tags.items() if not _is_reserved_tag(key)}
+
+    @property
+    def tags(self) -> dict[str, str]:
+        """
+        Return the prompt-level tags (from RegisteredModel).
+        """
+        return {key: value for key, value in self._prompt_tags.items() if not _is_reserved_tag(key)}
+
+    @property
+    def run_ids(self) -> list[str]:
+        """Get the run IDs associated with the prompt."""
+        run_tag = self._tags.get(PROMPT_ASSOCIATED_RUN_IDS_TAG_KEY)
+        if not run_tag:
+            return []
+        return run_tag.split(",")
+
+    @property
+    def uri(self) -> str:
+        """Return the URI of the prompt."""
+        return f"prompts:/{self.name}/{self.version}"
+
+    def format(self, allow_partial: bool = False, **kwargs) -> Union[Prompt, str]:
+        """
+        Format the template text with the given keyword arguments.
+        By default, it raises an error if there are missing variables. To format
+        the prompt text partially, set `allow_partial=True`.
+
+        Example:
+
+        .. code-block:: python
+
+            prompt = Prompt("my-prompt", 1, "Hello, {{title}} {{name}}!")
+            formatted = prompt.format(title="Ms", name="Alice")
+            print(formatted)
+            # Output: "Hello, Ms Alice!"
+
+            # Partial formatting
+            formatted = prompt.format(title="Ms", allow_partial=True)
+            print(formatted)
+            # Output: Prompt(name=my-prompt, version=1, template="Hello, Ms {{name}}!")
+
+
+        Args:
+            allow_partial: If True, allow partial formatting of the prompt text.
+                If False, raise an error if there are missing variables.
+            kwargs: Keyword arguments to replace the variables in the template.
+        """
+        input_keys = set(kwargs.keys())
+
+        template = self.template
+        for key, value in kwargs.items():
+            template = re.sub(r"\{\{\s*" + key + r"\s*\}\}", str(value), template)
+
+        if missing_keys := self.variables - input_keys:
+            if not allow_partial:
+                raise MlflowException.invalid_parameter_value(
+                    f"Missing variables: {missing_keys}. To partially format the prompt, "
+                    "set `allow_partial=True`."
+                )
+            else:
+                return Prompt(
+                    name=self.name,
+                    version=self.version,
+                    template=template,
+                    commit_message=self.commit_message,
+                    creation_timestamp=self.creation_timestamp,
+                    prompt_tags=self._prompt_tags,
+                    version_metadata=self.version_metadata,
+                    aliases=self.aliases,
+                )
+        return template
+
+    @classmethod
+    def from_model_version(
+        cls, model_version: ModelVersion, prompt_tags: Optional[dict[str, str]] = None
+    ) -> Prompt:
+        """
+        Create a Prompt object from a ModelVersion object.
+
+        Args:
+            model_version: The ModelVersion object to convert to a Prompt.
+            prompt_tags: The prompt-level tags (from RegisteredModel). Optional.
+        """
+        if IS_PROMPT_TAG_KEY not in model_version.tags:
+            raise MlflowException.invalid_parameter_value(
+                f"Name `{model_version.name}` is registered as a model, not a prompt. MLflow "
+                "does not allow registering a prompt with the same name as an existing model.",
+            )
+
+        if PROMPT_TEXT_TAG_KEY not in model_version.tags:
+            raise MlflowException.invalid_parameter_value(
+                f"Prompt `{model_version.name}` does not contain a prompt text"
+            )
+
+        return cls(
+            name=model_version.name,
+            version=model_version.version,
+            template=model_version.tags[PROMPT_TEXT_TAG_KEY],
+            commit_message=model_version.description,
+            creation_timestamp=model_version.creation_timestamp,
+            version_metadata=model_version.tags,
+            prompt_tags=prompt_tags,
+            aliases=model_version.aliases,
+        )

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/registered_model.py

@@ -0,0 +1,144 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.entities.model_registry.model_version import ModelVersion
+from mlflow.entities.model_registry.prompt import IS_PROMPT_TAG_KEY
+from mlflow.entities.model_registry.registered_model_alias import RegisteredModelAlias
+from mlflow.entities.model_registry.registered_model_tag import RegisteredModelTag
+from mlflow.protos.model_registry_pb2 import RegisteredModel as ProtoRegisteredModel
+from mlflow.protos.model_registry_pb2 import RegisteredModelAlias as ProtoRegisteredModelAlias
+from mlflow.protos.model_registry_pb2 import RegisteredModelTag as ProtoRegisteredModelTag
+
+
+class RegisteredModel(_ModelRegistryEntity):
+    """
+    MLflow entity for Registered Model.
+    """
+
+    def __init__(
+        self,
+        name,
+        creation_timestamp=None,
+        last_updated_timestamp=None,
+        description=None,
+        latest_versions=None,
+        tags=None,
+        aliases=None,
+    ):
+        # Constructor is called only from within the system by various backend stores.
+        super().__init__()
+        self._name = name
+        self._creation_time = creation_timestamp
+        self._last_updated_timestamp = last_updated_timestamp
+        self._description = description
+        self._latest_version = latest_versions
+        self._tags = {tag.key: tag.value for tag in (tags or [])}
+        self._aliases = {alias.alias: alias.version for alias in (aliases or [])}
+
+    @property
+    def name(self):
+        """String. Registered model name."""
+        return self._name
+
+    @name.setter
+    def name(self, new_name):
+        self._name = new_name
+
+    @property
+    def creation_timestamp(self):
+        """Integer. Model version creation timestamp (milliseconds since the Unix epoch)."""
+        return self._creation_time
+
+    @property
+    def last_updated_timestamp(self):
+        """Integer. Timestamp of last update for this model version (milliseconds since the Unix
+        epoch).
+        """
+        return self._last_updated_timestamp
+
+    @last_updated_timestamp.setter
+    def last_updated_timestamp(self, updated_timestamp):
+        self._last_updated_timestamp = updated_timestamp
+
+    @property
+    def description(self):
+        """String. Description"""
+        return self._description
+
+    @description.setter
+    def description(self, description):
+        self._description = description
+
+    @property
+    def latest_versions(self):
+        """List of the latest :py:class:`mlflow.entities.model_registry.ModelVersion` instances
+        for each stage.
+        """
+        return self._latest_version
+
+    @latest_versions.setter
+    def latest_versions(self, latest_versions):
+        self._latest_version = latest_versions
+
+    @property
+    def tags(self):
+        """Dictionary of tag key (string) -> tag value for the current registered model."""
+        # Remove the is_prompt tag as it should not be user-facing
+        return {k: v for k, v in self._tags.items() if k != IS_PROMPT_TAG_KEY}
+
+    @property
+    def aliases(self):
+        """Dictionary of aliases (string) -> version for the current registered model."""
+        return self._aliases
+
+    @classmethod
+    def _properties(cls):
+        # aggregate with base class properties since cls.__dict__ does not do it automatically
+        return sorted(cls._get_properties_helper())
+
+    def _add_tag(self, tag):
+        self._tags[tag.key] = tag.value
+
+    def _add_alias(self, alias):
+        self._aliases[alias.alias] = alias.version
+
+    # proto mappers
+    @classmethod
+    def from_proto(cls, proto):
+        # input: mlflow.protos.model_registry_pb2.RegisteredModel
+        # returns RegisteredModel entity
+        registered_model = cls(
+            proto.name,
+            proto.creation_timestamp,
+            proto.last_updated_timestamp,
+            proto.description,
+            [ModelVersion.from_proto(mvd) for mvd in proto.latest_versions],
+        )
+        for tag in proto.tags:
+            registered_model._add_tag(RegisteredModelTag.from_proto(tag))
+        for alias in proto.aliases:
+            registered_model._add_alias(RegisteredModelAlias.from_proto(alias))
+        return registered_model
+
+    def to_proto(self):
+        # returns mlflow.protos.model_registry_pb2.RegisteredModel
+        rmd = ProtoRegisteredModel()
+        rmd.name = self.name
+        if self.creation_timestamp is not None:
+            rmd.creation_timestamp = self.creation_timestamp
+        if self.last_updated_timestamp:
+            rmd.last_updated_timestamp = self.last_updated_timestamp
+        if self.description:
+            rmd.description = self.description
+        if self.latest_versions is not None:
+            rmd.latest_versions.extend(
+                [model_version.to_proto() for model_version in self.latest_versions]
+            )
+        rmd.tags.extend(
+            [ProtoRegisteredModelTag(key=key, value=value) for key, value in self._tags.items()]
+        )
+        rmd.aliases.extend(
+            [
+                ProtoRegisteredModelAlias(alias=alias, version=str(version))
+                for alias, version in self._aliases.items()
+            ]
+        )
+        return rmd

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/registered_model_alias.py

@@ -0,0 +1,35 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.protos.model_registry_pb2 import RegisteredModelAlias as ProtoRegisteredModelAlias
+
+
+class RegisteredModelAlias(_ModelRegistryEntity):
+    """Alias object associated with a registered model."""
+
+    def __init__(self, alias, version):
+        self._alias = alias
+        self._version = version
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def alias(self):
+        """String name of the alias."""
+        return self._alias
+
+    @property
+    def version(self):
+        """String model version number that the alias points to."""
+        return self._version
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.alias, proto.version)
+
+    def to_proto(self):
+        alias_proto = ProtoRegisteredModelAlias()
+        alias_proto.alias = self.alias
+        alias_proto.version = self.version
+        return alias_proto

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/registered_model_search.py

@@ -0,0 +1,25 @@
+from mlflow.entities.model_registry import RegisteredModel
+
+
+class RegisteredModelSearch(RegisteredModel):
+    def __init__(self, *args, **kwargs):
+        kwargs["tags"] = []
+        kwargs["aliases"] = []
+        super().__init__(*args, **kwargs)
+
+    def tags(self):
+        raise Exception(
+            "UC Registered Models gathered through search_registered_models do not have tags. "
+            "Please use get_registered_model to obtain an individual model's tags."
+        )
+
+    def aliases(self):
+        raise Exception(
+            "UC Registered Models gathered through search_registered_models do not have aliases. "
+            "Please use get_registered_model to obtain an individual model's aliases."
+        )
+
+    def __eq__(self, other):
+        if type(other) in {type(self), RegisteredModel}:
+            return self.__dict__ == other.__dict__
+        return False

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/registered_model_tag.py

@@ -0,0 +1,35 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.protos.model_registry_pb2 import RegisteredModelTag as ProtoRegisteredModelTag
+
+
+class RegisteredModelTag(_ModelRegistryEntity):
+    """Tag object associated with a registered model."""
+
+    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
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)
+
+    def to_proto(self):
+        tag = ProtoRegisteredModelTag()
+        tag.key = self.key
+        tag.value = self.value
+        return tag

venv/lib/python3.9/site-packages/mlflow/entities/multipart_upload.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from mlflow.protos.mlflow_artifacts_pb2 import (
+    CreateMultipartUpload as ProtoCreateMultipartUpload,
+)
+from mlflow.protos.mlflow_artifacts_pb2 import (
+    MultipartUploadCredential as ProtoMultipartUploadCredential,
+)
+
+
+@dataclass
+class MultipartUploadPart:
+    part_number: int
+    etag: str
+    url: Optional[str] = None
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            proto.part_number,
+            proto.etag or None,
+            proto.url or None,
+        )
+
+    def to_dict(self):
+        return {
+            "part_number": self.part_number,
+            "etag": self.etag,
+            "url": self.url,
+        }
+
+
+@dataclass
+class MultipartUploadCredential:
+    url: str
+    part_number: int
+    headers: dict[str, Any]
+
+    def to_proto(self):
+        credential = ProtoMultipartUploadCredential()
+        credential.url = self.url
+        credential.part_number = self.part_number
+        credential.headers.update(self.headers)
+        return credential
+
+    @classmethod
+    def from_dict(cls, dict_):
+        return cls(
+            url=dict_["url"],
+            part_number=dict_["part_number"],
+            headers=dict_.get("headers", {}),
+        )
+
+
+@dataclass
+class CreateMultipartUploadResponse:
+    upload_id: Optional[str]
+    credentials: list[MultipartUploadCredential]
+
+    def to_proto(self):
+        response = ProtoCreateMultipartUpload.Response()
+        if self.upload_id:
+            response.upload_id = self.upload_id
+        response.credentials.extend([credential.to_proto() for credential in self.credentials])
+        return response
+
+    @classmethod
+    def from_dict(cls, dict_):
+        credentials = [MultipartUploadCredential.from_dict(cred) for cred in dict_["credentials"]]
+        return cls(
+            upload_id=dict_.get("upload_id"),
+            credentials=credentials,
+        )

venv/lib/python3.9/site-packages/mlflow/entities/param.py

@@ -0,0 +1,49 @@
+import sys
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import Param as ProtoParam
+
+
+class Param(_MlflowObject):
+    """
+    Parameter object.
+    """
+
+    def __init__(self, key, value):
+        if "pyspark.ml" in sys.modules:
+            import pyspark.ml.param
+
+            if isinstance(key, pyspark.ml.param.Param):
+                key = key.name
+                value = str(value)
+        self._key = key
+        self._value = value
+
+    @property
+    def key(self):
+        """String key corresponding to the parameter name."""
+        return self._key
+
+    @property
+    def value(self):
+        """String value of the parameter."""
+        return self._value
+
+    def to_proto(self):
+        param = ProtoParam()
+        param.key = self.key
+        param.value = self.value
+        return param
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self._key == __o._key
+
+        return False
+
+    def __hash__(self):
+        return hash(self._key)

venv/lib/python3.9/site-packages/mlflow/entities/run.py

@@ -0,0 +1,77 @@
+from typing import Any, Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.run_data import RunData
+from mlflow.entities.run_info import RunInfo
+from mlflow.entities.run_inputs import RunInputs
+from mlflow.exceptions import MlflowException
+from mlflow.protos.service_pb2 import Run as ProtoRun
+
+
+class Run(_MlflowObject):
+    """
+    Run object.
+    """
+
+    def __init__(
+        self, run_info: RunInfo, run_data: RunData, run_inputs: Optional[RunInputs] = None
+    ) -> None:
+        if run_info is None:
+            raise MlflowException("run_info cannot be None")
+        self._info = run_info
+        self._data = run_data
+        self._inputs = run_inputs
+
+    @property
+    def info(self) -> RunInfo:
+        """
+        The run metadata, such as the run id, start time, and status.
+
+        :rtype: :py:class:`mlflow.entities.RunInfo`
+        """
+        return self._info
+
+    @property
+    def data(self) -> RunData:
+        """
+        The run data, including metrics, parameters, and tags.
+
+        :rtype: :py:class:`mlflow.entities.RunData`
+        """
+        return self._data
+
+    @property
+    def inputs(self) -> RunInputs:
+        """
+        The run inputs, including dataset inputs
+
+        :rtype: :py:class:`mlflow.entities.RunInputs`
+        """
+        return self._inputs
+
+    def to_proto(self):
+        run = ProtoRun()
+        run.info.MergeFrom(self.info.to_proto())
+        if self.data:
+            run.data.MergeFrom(self.data.to_proto())
+        if self.inputs:
+            run.inputs.MergeFrom(self.inputs.to_proto())
+        return run
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            RunInfo.from_proto(proto.info),
+            RunData.from_proto(proto.data),
+            RunInputs.from_proto(proto.inputs),
+        )
+
+    def to_dictionary(self) -> dict[Any, Any]:
+        run_dict = {
+            "info": dict(self.info),
+        }
+        if self.data:
+            run_dict["data"] = self.data.to_dictionary()
+        if self.inputs:
+            run_dict["inputs"] = self.inputs.to_dictionary()
+        return run_dict

venv/lib/python3.9/site-packages/mlflow/entities/run_data.py

@@ -0,0 +1,84 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.metric import Metric
+from mlflow.entities.param import Param
+from mlflow.entities.run_tag import RunTag
+from mlflow.protos.service_pb2 import Param as ProtoParam
+from mlflow.protos.service_pb2 import RunData as ProtoRunData
+from mlflow.protos.service_pb2 import RunTag as ProtoRunTag
+
+
+class RunData(_MlflowObject):
+    """
+    Run data (metrics and parameters).
+    """
+
+    def __init__(self, metrics=None, params=None, tags=None):
+        """Construct a new mlflow.entities.RunData instance.
+
+        Args:
+            metrics: List of mlflow.entities.Metric.
+            params: List of mlflow.entities.Param.
+            tags: List of mlflow.entities.RunTag.
+
+        """
+        # Maintain the original list of metrics so that we can easily convert it back to
+        # protobuf
+        self._metric_objs = metrics or []
+        self._metrics = {metric.key: metric.value for metric in self._metric_objs}
+        self._params = {param.key: param.value for param in (params or [])}
+        self._tags = {tag.key: tag.value for tag in (tags or [])}
+
+    @property
+    def metrics(self):
+        """
+        Dictionary of string key -> metric value for the current run.
+        For each metric key, the metric value with the latest timestamp is returned. In case there
+        are multiple values with the same latest timestamp, the maximum of these values is returned.
+        """
+        return self._metrics
+
+    @property
+    def params(self):
+        """Dictionary of param key (string) -> param value for the current run."""
+        return self._params
+
+    @property
+    def tags(self):
+        """Dictionary of tag key (string) -> tag value for the current run."""
+        return self._tags
+
+    def _add_metric(self, metric):
+        self._metrics[metric.key] = metric.value
+        self._metric_objs.append(metric)
+
+    def _add_param(self, param):
+        self._params[param.key] = param.value
+
+    def _add_tag(self, tag):
+        self._tags[tag.key] = tag.value
+
+    def to_proto(self):
+        run_data = ProtoRunData()
+        run_data.metrics.extend([m.to_proto() for m in self._metric_objs])
+        run_data.params.extend([ProtoParam(key=key, value=val) for key, val in self.params.items()])
+        run_data.tags.extend([ProtoRunTag(key=key, value=val) for key, val in self.tags.items()])
+        return run_data
+
+    def to_dictionary(self):
+        return {
+            "metrics": self.metrics,
+            "params": self.params,
+            "tags": self.tags,
+        }
+
+    @classmethod
+    def from_proto(cls, proto):
+        run_data = cls()
+        # iterate proto and add metrics, params, and tags
+        for proto_metric in proto.metrics:
+            run_data._add_metric(Metric.from_proto(proto_metric))
+        for proto_param in proto.params:
+            run_data._add_param(Param.from_proto(proto_param))
+        for proto_tag in proto.tags:
+            run_data._add_tag(RunTag.from_proto(proto_tag))
+        return run_data

venv/lib/python3.9/site-packages/mlflow/entities/run_info.py

@@ -0,0 +1,207 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.lifecycle_stage import LifecycleStage
+from mlflow.entities.run_status import RunStatus
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.service_pb2 import RunInfo as ProtoRunInfo
+
+
+def check_run_is_active(run_info):
+    if run_info.lifecycle_stage != LifecycleStage.ACTIVE:
+        raise MlflowException(
+            f"The run {run_info.run_id} must be in 'active' lifecycle_stage.",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+
+class searchable_attribute(property):
+    # Wrapper class over property to designate some of the properties as searchable
+    # run attributes
+    pass
+
+
+class orderable_attribute(property):
+    # Wrapper class over property to designate some of the properties as orderable
+    # run attributes
+    pass
+
+
+class RunInfo(_MlflowObject):
+    """
+    Metadata about a run.
+    """
+
+    def __init__(
+        self,
+        run_uuid,
+        experiment_id,
+        user_id,
+        status,
+        start_time,
+        end_time,
+        lifecycle_stage,
+        artifact_uri=None,
+        run_id=None,
+        run_name=None,
+    ):
+        if experiment_id is None:
+            raise Exception("experiment_id cannot be None")
+        if user_id is None:
+            raise Exception("user_id cannot be None")
+        if status is None:
+            raise Exception("status cannot be None")
+        if start_time is None:
+            raise Exception("start_time cannot be None")
+        actual_run_id = run_id or run_uuid
+        if actual_run_id is None:
+            raise Exception("run_id and run_uuid cannot both be None")
+        self._run_uuid = actual_run_id
+        self._run_id = actual_run_id
+        self._experiment_id = experiment_id
+        self._user_id = user_id
+        self._status = status
+        self._start_time = start_time
+        self._end_time = end_time
+        self._lifecycle_stage = lifecycle_stage
+        self._artifact_uri = artifact_uri
+        self._run_name = run_name
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            # TODO deep equality here?
+            return self.__dict__ == other.__dict__
+        return False
+
+    def _copy_with_overrides(self, status=None, end_time=None, lifecycle_stage=None, run_name=None):
+        """A copy of the RunInfo with certain attributes modified."""
+        proto = self.to_proto()
+        if status:
+            proto.status = status
+        if end_time:
+            proto.end_time = end_time
+        if lifecycle_stage:
+            proto.lifecycle_stage = lifecycle_stage
+        if run_name:
+            proto.run_name = run_name
+        return RunInfo.from_proto(proto)
+
+    @property
+    def run_uuid(self):
+        """[Deprecated, use run_id instead] String containing run UUID."""
+        return self._run_uuid
+
+    @searchable_attribute
+    def run_id(self):
+        """String containing run id."""
+        return self._run_id
+
+    @property
+    def experiment_id(self):
+        """String ID of the experiment for the current run."""
+        return self._experiment_id
+
+    @searchable_attribute
+    def run_name(self):
+        """String containing run name."""
+        return self._run_name
+
+    def _set_run_name(self, new_name):
+        self._run_name = new_name
+
+    @searchable_attribute
+    def user_id(self):
+        """String ID of the user who initiated this run."""
+        return self._user_id
+
+    @searchable_attribute
+    def status(self):
+        """
+        One of the values in :py:class:`mlflow.entities.RunStatus`
+        describing the status of the run.
+        """
+        return self._status
+
+    @searchable_attribute
+    def start_time(self):
+        """Start time of the run, in number of milliseconds since the UNIX epoch."""
+        return self._start_time
+
+    @searchable_attribute
+    def end_time(self):
+        """End time of the run, in number of milliseconds since the UNIX epoch."""
+        return self._end_time
+
+    @searchable_attribute
+    def artifact_uri(self):
+        """String root artifact URI of the run."""
+        return self._artifact_uri
+
+    @property
+    def lifecycle_stage(self):
+        """
+        One of the values in :py:class:`mlflow.entities.lifecycle_stage.LifecycleStage`
+        describing the lifecycle stage of the run.
+        """
+        return self._lifecycle_stage
+
+    def to_proto(self):
+        proto = ProtoRunInfo()
+        proto.run_uuid = self.run_uuid
+        proto.run_id = self.run_id
+        if self.run_name is not None:
+            proto.run_name = self.run_name
+        proto.experiment_id = self.experiment_id
+        proto.user_id = self.user_id
+        proto.status = RunStatus.from_string(self.status)
+        proto.start_time = self.start_time
+        if self.end_time:
+            proto.end_time = self.end_time
+        if self.artifact_uri:
+            proto.artifact_uri = self.artifact_uri
+        proto.lifecycle_stage = self.lifecycle_stage
+        return proto
+
+    @classmethod
+    def from_proto(cls, proto):
+        end_time = proto.end_time
+        # The proto2 default scalar value of zero indicates that the run's end time is absent.
+        # An absent end time is represented with a NoneType in the `RunInfo` class
+        if end_time == 0:
+            end_time = None
+        return cls(
+            run_uuid=proto.run_uuid,
+            run_id=proto.run_id,
+            run_name=proto.run_name,
+            experiment_id=proto.experiment_id,
+            user_id=proto.user_id,
+            status=RunStatus.to_string(proto.status),
+            start_time=proto.start_time,
+            end_time=end_time,
+            lifecycle_stage=proto.lifecycle_stage,
+            artifact_uri=proto.artifact_uri,
+        )
+
+    @classmethod
+    def get_searchable_attributes(cls):
+        return sorted(
+            [p for p in cls.__dict__ if isinstance(getattr(cls, p), searchable_attribute)]
+        )
+
+    @classmethod
+    def get_orderable_attributes(cls):
+        # Note that all searchable attributes are also orderable.
+        return sorted(
+            [
+                p
+                for p in cls.__dict__
+                if isinstance(getattr(cls, p), (searchable_attribute, orderable_attribute))
+            ]
+        )
+
+    @classmethod
+    def from_dictionary(cls, the_dict) -> "RunInfo":
+        # To support loading runs created in 3.x, backfill `run_uuid` if it is not present
+        extra = {}
+        if "run_uuid" not in the_dict and (run_id := the_dict.get("run_id")) is not None:
+            extra["run_uuid"] = run_id
+        return super().from_dictionary(the_dict | extra)

venv/lib/python3.9/site-packages/mlflow/entities/run_inputs.py

@@ -0,0 +1,41 @@
+from typing import Any
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.dataset_input import DatasetInput
+from mlflow.protos.service_pb2 import RunInputs as ProtoRunInputs
+
+
+class RunInputs(_MlflowObject):
+    """RunInputs object."""
+
+    def __init__(self, dataset_inputs: list[DatasetInput]) -> None:
+        self._dataset_inputs = dataset_inputs
+
+    def __eq__(self, other: _MlflowObject) -> bool:
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def dataset_inputs(self) -> list[DatasetInput]:
+        """Array of dataset inputs."""
+        return self._dataset_inputs
+
+    def to_proto(self):
+        run_inputs = ProtoRunInputs()
+        run_inputs.dataset_inputs.extend(
+            [dataset_input.to_proto() for dataset_input in self.dataset_inputs]
+        )
+        return run_inputs
+
+    def to_dictionary(self) -> dict[str, Any]:
+        return {
+            "dataset_inputs": [d.to_dictionary() for d in self.dataset_inputs],
+        }
+
+    @classmethod
+    def from_proto(cls, proto):
+        dataset_inputs = [
+            DatasetInput.from_proto(dataset_input) for dataset_input in proto.dataset_inputs
+        ]
+        return cls(dataset_inputs)

venv/lib/python3.9/site-packages/mlflow/entities/run_status.py

@@ -0,0 +1,41 @@
+from mlflow.protos.service_pb2 import RunStatus as ProtoRunStatus
+
+
+class RunStatus:
+    """Enum for status of an :py:class:`mlflow.entities.Run`."""
+
+    RUNNING = ProtoRunStatus.Value("RUNNING")
+    SCHEDULED = ProtoRunStatus.Value("SCHEDULED")
+    FINISHED = ProtoRunStatus.Value("FINISHED")
+    FAILED = ProtoRunStatus.Value("FAILED")
+    KILLED = ProtoRunStatus.Value("KILLED")
+
+    _STRING_TO_STATUS = {k: ProtoRunStatus.Value(k) for k in ProtoRunStatus.keys()}
+    _STATUS_TO_STRING = {value: key for key, value in _STRING_TO_STATUS.items()}
+    _TERMINATED_STATUSES = {FINISHED, FAILED, KILLED}
+
+    @staticmethod
+    def from_string(status_str):
+        if status_str not in RunStatus._STRING_TO_STATUS:
+            raise Exception(
+                f"Could not get run status corresponding to string {status_str}. Valid run "
+                f"status strings: {list(RunStatus._STRING_TO_STATUS.keys())}"
+            )
+        return RunStatus._STRING_TO_STATUS[status_str]
+
+    @staticmethod
+    def to_string(status):
+        if status not in RunStatus._STATUS_TO_STRING:
+            raise Exception(
+                f"Could not get string corresponding to run status {status}. Valid run "
+                f"statuses: {list(RunStatus._STATUS_TO_STRING.keys())}"
+            )
+        return RunStatus._STATUS_TO_STRING[status]
+
+    @staticmethod
+    def is_terminated(status):
+        return status in RunStatus._TERMINATED_STATUSES
+
+    @staticmethod
+    def all_status():
+        return list(RunStatus._STATUS_TO_STRING.keys())

venv/lib/python3.9/site-packages/mlflow/entities/run_tag.py

@@ -0,0 +1,36 @@
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import RunTag as ProtoRunTag
+
+
+class RunTag(_MlflowObject):
+    """Tag object associated with a run."""
+
+    def __init__(self, key, value):
+        self._key = key
+        self._value = value
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            # TODO deep equality here?
+            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 = ProtoRunTag()
+        param.key = self.key
+        param.value = self.value
+        return param
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)

venv/lib/python3.9/site-packages/mlflow/entities/source_type.py

@@ -0,0 +1,32 @@
+class SourceType:
+    """Enum for originating source of a :py:class:`mlflow.entities.Run`."""
+
+    NOTEBOOK, JOB, PROJECT, LOCAL, UNKNOWN, RECIPE = range(1, 7)
+
+    _STRING_TO_SOURCETYPE = {
+        "NOTEBOOK": NOTEBOOK,
+        "JOB": JOB,
+        "PROJECT": PROJECT,
+        "LOCAL": LOCAL,
+        "RECIPE": RECIPE,
+        "UNKNOWN": UNKNOWN,
+    }
+    SOURCETYPE_TO_STRING = {value: key for key, value in _STRING_TO_SOURCETYPE.items()}
+
+    @staticmethod
+    def from_string(status_str):
+        if status_str not in SourceType._STRING_TO_SOURCETYPE:
+            raise Exception(
+                f"Could not get run status corresponding to string {status_str}. Valid run "
+                f"status strings: {list(SourceType._STRING_TO_SOURCETYPE.keys())}"
+            )
+        return SourceType._STRING_TO_SOURCETYPE[status_str]
+
+    @staticmethod
+    def to_string(status):
+        if status not in SourceType.SOURCETYPE_TO_STRING:
+            raise Exception(
+                f"Could not get string corresponding to run status {status}. Valid run "
+                f"statuses: {list(SourceType.SOURCETYPE_TO_STRING.keys())}"
+            )
+        return SourceType.SOURCETYPE_TO_STRING[status]

venv/lib/python3.9/site-packages/mlflow/entities/span.py

@@ -0,0 +1,678 @@
+import base64
+import json
+import logging
+from dataclasses import asdict
+from functools import lru_cache
+from typing import Any, Optional, Union
+
+from opentelemetry.sdk.trace import Event as OTelEvent
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
+from opentelemetry.trace import Span as OTelSpan
+
+import mlflow
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatus, SpanStatusCode
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.databricks_trace_server_pb2 import Span as ProtoSpan
+from mlflow.tracing.constant import SpanAttributeKey
+from mlflow.tracing.utils import (
+    TraceJSONEncoder,
+    build_otel_context,
+    decode_id,
+    encode_span_id,
+    encode_trace_id,
+)
+from mlflow.utils.proto_json_utils import set_pb_value
+
+_logger = logging.getLogger(__name__)
+
+
+# Not using enum as we want to allow custom span type string.
+class SpanType:
+    """
+    Predefined set of span types.
+    """
+
+    LLM = "LLM"
+    CHAIN = "CHAIN"
+    AGENT = "AGENT"
+    TOOL = "TOOL"
+    CHAT_MODEL = "CHAT_MODEL"
+    RETRIEVER = "RETRIEVER"
+    PARSER = "PARSER"
+    EMBEDDING = "EMBEDDING"
+    RERANKER = "RERANKER"
+    UNKNOWN = "UNKNOWN"
+
+
+def create_mlflow_span(
+    otel_span: Any, request_id: str, span_type: Optional[str] = None
+) -> Union["Span", "LiveSpan", "NoOpSpan"]:
+    """
+    Factory function to create a span object.
+
+    When creating a MLflow span object from the OpenTelemetry span, the factory function
+    should always be used to ensure the correct span object is created.
+    """
+    if not otel_span or isinstance(otel_span, NonRecordingSpan):
+        return NoOpSpan()
+
+    if isinstance(otel_span, OTelSpan):
+        return LiveSpan(otel_span, request_id, span_type)
+
+    if isinstance(otel_span, OTelReadableSpan):
+        return Span(otel_span)
+
+    raise MlflowException(
+        "The `otel_span` argument must be an instance of one of valid "
+        f"OpenTelemetry span classes, but got {type(otel_span)}.",
+        INVALID_PARAMETER_VALUE,
+    )
+
+
+class Span:
+    """
+    A span object. A span represents a unit of work or operation and is the building
+    block of Traces.
+
+    This Span class represents immutable span data that is already finished and persisted.
+    The "live" span that is being created and updated during the application runtime is
+    represented by the :py:class:`LiveSpan <mlflow.entities.LiveSpan>` subclass.
+    """
+
+    def __init__(self, otel_span: OTelReadableSpan):
+        if not isinstance(otel_span, OTelReadableSpan):
+            raise MlflowException(
+                "The `otel_span` argument for the Span class must be an instance of ReadableSpan, "
+                f"but got {type(otel_span)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        self._span = otel_span
+        # Since the span is immutable, we can cache the attributes to avoid the redundant
+        # deserialization of the attribute values.
+        self._attributes = _CachedSpanAttributesRegistry(otel_span)
+
+    @property
+    @lru_cache(maxsize=1)
+    def request_id(self) -> str:
+        """
+        The request ID of the span, a unique identifier for the trace it belongs to.
+        Request ID is equivalent to the trace ID in OpenTelemetry, but generated
+        differently by the tracing backend.
+        """
+        return self.get_attribute(SpanAttributeKey.REQUEST_ID)
+
+    @property
+    def span_id(self) -> str:
+        """The ID of the span. This is only unique within a trace."""
+        return encode_span_id(self._span.context.span_id)
+
+    @property
+    def name(self) -> str:
+        """The name of the span."""
+        return self._span.name
+
+    @property
+    def start_time_ns(self) -> int:
+        """The start time of the span in nanosecond."""
+        return self._span._start_time
+
+    @property
+    def end_time_ns(self) -> Optional[int]:
+        """The end time of the span in nanosecond."""
+        return self._span._end_time
+
+    @property
+    def parent_id(self) -> Optional[str]:
+        """The span ID of the parent span."""
+        if self._span.parent is None:
+            return None
+        return encode_span_id(self._span.parent.span_id)
+
+    @property
+    def status(self) -> SpanStatus:
+        """The status of the span."""
+        return SpanStatus.from_otel_status(self._span.status)
+
+    @property
+    def inputs(self) -> Any:
+        """The input values of the span."""
+        return self.get_attribute(SpanAttributeKey.INPUTS)
+
+    @property
+    def outputs(self) -> Any:
+        """The output values of the span."""
+        return self.get_attribute(SpanAttributeKey.OUTPUTS)
+
+    @property
+    def span_type(self) -> str:
+        """The type of the span."""
+        return self.get_attribute(SpanAttributeKey.SPAN_TYPE)
+
+    @property
+    def _trace_id(self) -> str:
+        """
+        The OpenTelemetry trace ID of the span. Note that this should not be exposed to
+        the user, instead, use request_id as an unique identifier for a trace.
+        """
+        return encode_trace_id(self._span.context.trace_id)
+
+    @property
+    def attributes(self) -> dict[str, Any]:
+        """
+        Get all attributes of the span.
+
+        Returns:
+            A dictionary of all attributes of the span.
+        """
+        return self._attributes.get_all()
+
+    @property
+    def events(self) -> list[SpanEvent]:
+        """
+        Get all events of the span.
+
+        Returns:
+            A list of all events of the span.
+        """
+        return [
+            SpanEvent(
+                name=event.name,
+                timestamp=event.timestamp,
+                # Convert from OpenTelemetry's BoundedAttributes class to a simple dict
+                # to avoid the serialization issue due to having a lock object.
+                attributes=dict(event.attributes),
+            )
+            for event in self._span.events
+        ]
+
+    def __repr__(self):
+        return (
+            f"{type(self).__name__}(name={self.name!r}, request_id={self.request_id!r}, "
+            f"span_id={self.span_id!r}, parent_id={self.parent_id!r})"
+        )
+
+    def get_attribute(self, key: str) -> Optional[Any]:
+        """
+        Get a single attribute value from the span.
+
+        Args:
+            key: The key of the attribute to get.
+
+        Returns:
+            The value of the attribute if it exists, otherwise None.
+        """
+        return self._attributes.get(key)
+
+    def to_dict(self):
+        # NB: OpenTelemetry Span has to_json() method, but it will write many fields that
+        #  we don't use e.g. links, kind, resource, trace_state, etc. So we manually
+        #  cherry-pick the fields we need here.
+        return {
+            "name": self.name,
+            "context": {
+                "span_id": self.span_id,
+                "trace_id": self._trace_id,
+            },
+            "parent_id": self.parent_id,
+            "start_time": self.start_time_ns,
+            "end_time": self.end_time_ns,
+            "status_code": self.status.status_code.value,
+            "status_message": self.status.description,
+            "attributes": dict(self._span.attributes),
+            "events": [asdict(event) for event in self.events],
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Span":
+        """
+        Create a Span object from the given dictionary.
+        """
+        if trace_id := _get_trace_id_if_v3_format(data):
+            raise MlflowException(
+                f"Failed to load trace with ID {trace_id}. It was created with MLflow 3.x, "
+                f"but you're using MLflow {mlflow.__version__}, which cannot load traces created "
+                "with MLflow 3.x. To properly search and view traces from both MLflow 2.x and 3.x, "
+                "run this command to upgrade MLflow and restart your python process:\n"
+                "```\n"
+                "pip install 'mlflow>=3.0.0'\n"
+                "```"
+            )
+
+        try:
+            request_id = data.get("attributes", {}).get(SpanAttributeKey.REQUEST_ID)
+            if not request_id:
+                raise MlflowException(
+                    f"The {SpanAttributeKey.REQUEST_ID} attribute is empty or missing.",
+                    INVALID_PARAMETER_VALUE,
+                )
+
+            trace_id = decode_id(data["context"]["trace_id"])
+            span_id = decode_id(data["context"]["span_id"])
+            parent_id = decode_id(data["parent_id"]) if data["parent_id"] else None
+
+            otel_span = OTelReadableSpan(
+                name=data["name"],
+                context=build_otel_context(trace_id, span_id),
+                parent=build_otel_context(trace_id, parent_id) if parent_id else None,
+                start_time=data["start_time"],
+                end_time=data["end_time"],
+                attributes=data["attributes"],
+                status=SpanStatus(data["status_code"], data["status_message"]).to_otel_status(),
+                events=[
+                    OTelEvent(
+                        name=event["name"],
+                        timestamp=event["timestamp"],
+                        attributes=event["attributes"],
+                    )
+                    for event in data["events"]
+                ],
+            )
+            return cls(otel_span)
+        except Exception as e:
+            raise MlflowException(
+                "Failed to create a Span object from the given dictionary",
+                INVALID_PARAMETER_VALUE,
+            ) from e
+
+    def to_proto(self):
+        """Convert into OTLP compatible proto object to sent to the Databricks Trace Server."""
+        otel_status = self._span.status
+        status = ProtoSpan.Status(
+            code=otel_status.status_code.value,
+            message=otel_status.description,
+        )
+        parent = _encode_span_id_to_byte(self._span.parent.span_id) if self._span.parent else b""
+
+        proto = ProtoSpan(
+            trace_id=_encode_trace_id_to_byte(self._span.context.trace_id),
+            span_id=_encode_span_id_to_byte(self._span.context.span_id),
+            trace_state=self._span.context.trace_state or "",
+            parent_span_id=parent,
+            name=self.name,
+            start_time_unix_nano=self._span.start_time,
+            end_time_unix_nano=self._span.end_time,
+            events=[event.to_proto() for event in self.events],
+            status=status,
+        )
+
+        # Trace server's proto uses map<string, google.protobuf.Value> for attributes
+        for key, value in self._span.attributes.items():
+            set_pb_value(proto.attributes[key], value)
+
+        return proto
+
+
+def _encode_span_id_to_byte(span_id: Optional[int]) -> bytes:
+    # https://github.com/open-telemetry/opentelemetry-python/blob/e01fa0c77a7be0af77d008a888c2b6a707b05c3d/exporter/opentelemetry-exporter-otlp-proto-common/src/opentelemetry/exporter/otlp/proto/common/_internal/__init__.py#L131
+    return span_id.to_bytes(length=8, byteorder="big", signed=False)
+
+
+def _encode_trace_id_to_byte(trace_id: int) -> bytes:
+    # https://github.com/open-telemetry/opentelemetry-python/blob/e01fa0c77a7be0af77d008a888c2b6a707b05c3d/exporter/opentelemetry-exporter-otlp-proto-common/src/opentelemetry/exporter/otlp/proto/common/_internal/__init__.py#L135
+    return trace_id.to_bytes(length=16, byteorder="big", signed=False)
+
+
+def _get_trace_id_if_v3_format(data: dict[str, Any]) -> Optional[str]:
+    """Return Trace ID if the span dictionary is in the V3 span format. Otherwise None."""
+    if "context" not in data and "trace_id" in data:
+        # Try to decode the trace ID as the V3 format (base64 encoded 16 bytes)
+        # If decoding fails, the trace is not in the V3 format
+        # Ref: https://github.com/mlflow/mlflow/blame/08b3bd9402293039cd685a0dfa23d7f7138c93f3/mlflow/entities/span.py#L339-L342
+        try:
+            trace_id_bytes = base64.b64decode(data["trace_id"])
+            otel_trace_id = int.from_bytes(trace_id_bytes, byteorder="big", signed=False)
+            return "tr-" + encode_trace_id(otel_trace_id)
+        except Exception:
+            pass
+
+
+class LiveSpan(Span):
+    """
+    A "live" version of the :py:class:`Span <mlflow.entities.Span>` class.
+
+    The live spans are those being created and updated during the application runtime.
+    When users start a new span using the tracing APIs within their code, this live span
+    object is returned to get and set the span attributes, status, events, and etc.
+    """
+
+    def __init__(
+        self,
+        otel_span: OTelSpan,
+        request_id: str,
+        span_type: str = SpanType.UNKNOWN,
+    ):
+        """
+        The `otel_span` argument takes an instance of OpenTelemetry Span class, which is
+        indeed a subclass of ReadableSpan. Thanks to this, the getter methods of the Span
+        class can be reused without any modification.
+
+        Note that the constructor doesn't call the super().__init__ method, because the Span
+        initialization logic is a bit different from the immutable span.
+        """
+        if not isinstance(otel_span, OTelReadableSpan):
+            raise MlflowException(
+                "The `otel_span` argument for the LiveSpan class must be an instance of "
+                f"trace.Span, but got {type(otel_span)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        self._span = otel_span
+        self._attributes = _SpanAttributesRegistry(otel_span)
+        self._attributes.set(SpanAttributeKey.REQUEST_ID, request_id)
+        self._attributes.set(SpanAttributeKey.SPAN_TYPE, span_type)
+
+    def set_span_type(self, span_type: str):
+        """Set the type of the span."""
+        self.set_attribute(SpanAttributeKey.SPAN_TYPE, span_type)
+
+    def set_inputs(self, inputs: Any):
+        """Set the input values to the span."""
+        self.set_attribute(SpanAttributeKey.INPUTS, inputs)
+
+    def set_outputs(self, outputs: Any):
+        """Set the output values to the span."""
+        self.set_attribute(SpanAttributeKey.OUTPUTS, outputs)
+
+    def set_attributes(self, attributes: dict[str, Any]):
+        """
+        Set the attributes to the span. The attributes must be a dictionary of key-value pairs.
+        This method is additive, i.e. it will add new attributes to the existing ones. If an
+        attribute with the same key already exists, it will be overwritten.
+        """
+        if not isinstance(attributes, dict):
+            _logger.warning(
+                f"Attributes must be a dictionary, but got {type(attributes)}. Skipping."
+            )
+            return
+
+        for key, value in attributes.items():
+            self.set_attribute(key, value)
+
+    def set_attribute(self, key: str, value: Any):
+        """Set a single attribute to the span."""
+        self._attributes.set(key, value)
+
+    def set_status(self, status: Union[SpanStatusCode, str]):
+        """
+        Set the status of the span.
+
+        Args:
+            status: The status of the span. This can be a
+                :py:class:`SpanStatus <mlflow.entities.SpanStatus>` object or a string representing
+                of the status code defined in
+                :py:class:`SpanStatusCode <mlflow.entities.SpanStatusCode>`
+                e.g. ``"OK"``, ``"ERROR"``.
+        """
+        if isinstance(status, str):
+            status = SpanStatus(status)
+
+        # NB: We need to set the OpenTelemetry native StatusCode, because span's set_status
+        #     method only accepts a StatusCode enum in their definition.
+        #     https://github.com/open-telemetry/opentelemetry-python/blob/8ed71b15fb8fc9534529da8ce4a21e686248a8f3/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py#L949
+        #     Working around this is possible, but requires some hack to handle automatic status
+        #     propagation mechanism, so here we just use the native object that meets our
+        #     current requirements at least. Nevertheless, declaring the new class extending
+        #     the OpenTelemetry Status class so users code doesn't have to import the OTel's
+        #     StatusCode object, which makes future migration easier.
+        self._span.set_status(status.to_otel_status())
+
+    def add_event(self, event: SpanEvent):
+        """
+        Add an event to the span.
+
+        Args:
+            event: The event to add to the span. This should be a
+                :py:class:`SpanEvent <mlflow.entities.SpanEvent>` object.
+        """
+        self._span.add_event(event.name, event.attributes, event.timestamp)
+
+    def end(self, end_time: Optional[int] = None):
+        """
+        End the span. This is a thin wrapper around the OpenTelemetry's end method but just
+        to handle the status update.
+
+        This method should not be called directly by the user, only by called via fluent APIs
+        context exit or by MlflowClient APIs.
+
+        :meta private:
+        """
+        # NB: In OpenTelemetry, status code remains UNSET if not explicitly set
+        # by the user. However, there is not way to set the status when using
+        # @mlflow.trace decorator. Therefore, we just automatically set the status
+        # to OK if it is not ERROR.
+        if self.status.status_code != SpanStatusCode.ERROR:
+            self.set_status(SpanStatus(SpanStatusCode.OK))
+
+        self._span.end(end_time=end_time)
+
+    def from_dict(cls, data: dict[str, Any]) -> "Span":
+        raise NotImplementedError("The `from_dict` method is not supported for the LiveSpan class.")
+
+    def to_immutable_span(self) -> "Span":
+        """
+        Downcast the live span object to the immutable span.
+
+        :meta private:
+        """
+        # All state of the live span is already persisted in the OpenTelemetry span object.
+        return Span(self._span)
+
+    @classmethod
+    def from_immutable_span(
+        cls,
+        span: Span,
+        parent_span_id: Optional[str] = None,
+        request_id: Optional[str] = None,
+        trace_id: Optional[str] = None,
+    ) -> "LiveSpan":
+        """
+        Create a new LiveSpan object from the given immutable span by
+        cloning the underlying OpenTelemetry span within current context.
+
+        This is particularly useful when we merging a remote trace into the current trace.
+        We cannot merge the remote trace directly, because it is already stored as an immutable
+        span, meaning that we cannot update metadata like request ID, trace ID, parent span ID,
+        which are necessary for merging the trace.
+
+        Args:
+            span: The immutable span object to clone.
+            parent_span_id: The parent span ID of the new span.
+                If it is None, the span will be created as a root span.
+            request_id: The request ID to be set on the new span. Specify this if you want to
+                create the new span with a different request ID from the original span.
+            trace_id: The trace ID of the new span in hex encoded format. Specify this if you
+                want to create the new span with a different trace ID from the original span
+
+        Returns:
+            The new LiveSpan object with the same state as the original span.
+
+        :meta private:
+        """
+        from mlflow.tracing.trace_manager import InMemoryTraceManager
+
+        trace_manager = InMemoryTraceManager.get_instance()
+        request_id = request_id or span.request_id
+        parent_span = trace_manager.get_span_from_id(request_id, parent_span_id)
+
+        # Create a new span with the same name, parent, and start time
+        otel_span = mlflow.tracing.provider.start_detached_span(
+            name=span.name,
+            parent=parent_span._span if parent_span else None,
+            start_time_ns=span.start_time_ns,
+        )
+        # otel_span._span_processor = span._span._span_processor
+        clone_span = LiveSpan(otel_span, request_id, span.span_type)
+
+        # Copy all the attributes, inputs, outputs, and events from the original span
+        clone_span.set_status(span.status)
+        clone_span.set_attributes(
+            {k: v for k, v in span.attributes.items() if k != SpanAttributeKey.REQUEST_ID}
+        )
+        clone_span.set_inputs(span.inputs)
+        clone_span.set_outputs(span.outputs)
+        for event in span.events:
+            clone_span.add_event(event)
+
+        # Update trace ID and span ID
+        context = span._span.get_span_context()
+        clone_span._span._context = SpanContext(
+            # Override trace_id if provided, otherwise use the original trace ID
+            trace_id=decode_id(trace_id) or context.trace_id,
+            span_id=context.span_id,
+            is_remote=context.is_remote,
+            # Override trace flag as if it is sampled within current context.
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+        )
+
+        # Mark the span completed with the original end time
+        clone_span.end(end_time=span.end_time_ns)
+        return clone_span
+
+
+NO_OP_SPAN_REQUEST_ID = "MLFLOW_NO_OP_SPAN_REQUEST_ID"
+
+
+class NoOpSpan(Span):
+    """
+    No-op implementation of the Span interface.
+
+    This instance should be returned from the mlflow.start_span context manager when span
+    creation fails. This class should have exactly the same interface as the Span so that
+    user's setter calls do not raise runtime errors.
+
+    E.g.
+
+    .. code-block:: python
+
+        with mlflow.start_span("span_name") as span:
+            # Even if the span creation fails, the following calls should pass.
+            span.set_inputs({"x": 1})
+            # Do something
+
+    """
+
+    def __init__(self):
+        self._span = NonRecordingSpan(context=None)
+        self._attributes = {}
+
+    @property
+    def request_id(self):
+        """
+        No-op span returns a special request ID to distinguish it from the real spans.
+        """
+        return NO_OP_SPAN_REQUEST_ID
+
+    @property
+    def span_id(self):
+        return None
+
+    @property
+    def name(self):
+        return None
+
+    @property
+    def start_time_ns(self):
+        return None
+
+    @property
+    def end_time_ns(self):
+        return None
+
+    @property
+    def context(self):
+        return None
+
+    @property
+    def parent_id(self):
+        return None
+
+    @property
+    def status(self):
+        return None
+
+    @property
+    def _trace_id(self):
+        return None
+
+    def set_inputs(self, inputs: dict[str, Any]):
+        pass
+
+    def set_outputs(self, outputs: dict[str, Any]):
+        pass
+
+    def set_attributes(self, attributes: dict[str, Any]):
+        pass
+
+    def set_attribute(self, key: str, value: Any):
+        pass
+
+    def set_status(self, status: SpanStatus):
+        pass
+
+    def add_event(self, event: SpanEvent):
+        pass
+
+    def end(self):
+        pass
+
+
+class _SpanAttributesRegistry:
+    """
+    A utility class to manage the span attributes.
+
+    In MLflow users can add arbitrary key-value pairs to the span attributes, however,
+    OpenTelemetry only allows a limited set of types to be stored in the attribute values.
+    Therefore, we serialize all values into JSON string before storing them in the span.
+    This class provides simple getter and setter methods to interact with the span attributes
+    without worrying about the serde process.
+    """
+
+    def __init__(self, otel_span: OTelSpan):
+        self._span = otel_span
+
+    def get_all(self) -> dict[str, Any]:
+        return {key: self.get(key) for key in self._span.attributes.keys()}
+
+    def get(self, key: str):
+        serialized_value = self._span.attributes.get(key)
+        if serialized_value:
+            try:
+                return json.loads(serialized_value)
+            except Exception as e:
+                _logger.warning(
+                    f"Failed to get value for key {key}, make sure you set the attribute "
+                    f"on mlflow Span class instead of directly to the OpenTelemetry span. {e}"
+                )
+
+    def set(self, key: str, value: Any):
+        if not isinstance(key, str):
+            _logger.warning(f"Attribute key must be a string, but got {type(key)}. Skipping.")
+            return
+
+        # NB: OpenTelemetry attribute can store not only string but also a few primitives like
+        #   int, float, bool, and list of them. However, we serialize all into JSON string here
+        #   for the simplicity in deserialization process.
+        self._span.set_attribute(key, json.dumps(value, cls=TraceJSONEncoder, ensure_ascii=False))
+
+
+class _CachedSpanAttributesRegistry(_SpanAttributesRegistry):
+    """
+    A cache-enabled version of the SpanAttributesRegistry.
+
+    The caching helps to avoid the redundant deserialization of the attribute, however, it does
+    not handle the value change well. Therefore, this class should only be used for the persisted
+    spans that are immutable, and thus implemented as a subclass of _SpanAttributesRegistry.
+    """
+
+    @lru_cache(maxsize=128)
+    def get(self, key: str):
+        return super().get(key)
+
+    def set(self, key: str, value: Any):
+        raise MlflowException(
+            "The attributes of the immutable span must not be updated.", INVALID_PARAMETER_VALUE
+        )

venv/lib/python3.9/site-packages/mlflow/entities/span_event.py

@@ -0,0 +1,98 @@
+import json
+import sys
+import time
+import traceback
+from dataclasses import dataclass, field
+from datetime import datetime
+
+from opentelemetry.util.types import AttributeValue
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.databricks_trace_server_pb2 import Span as ProtoSpan
+from mlflow.utils.proto_json_utils import set_pb_value
+
+
+@dataclass
+class SpanEvent(_MlflowObject):
+    """
+    An event that records a specific occurrences or moments in time
+    during a span, such as an exception being thrown. Compatible with OpenTelemetry.
+
+    Args:
+        name: Name of the event.
+        timestamp:  The exact time the event occurred, measured in microseconds.
+            If not provided, the current time will be used.
+        attributes: A collection of key-value pairs representing detailed
+            attributes of the event, such as the exception stack trace.
+            Attributes value must be one of ``[str, int, float, bool, bytes]``
+            or a sequence of these types.
+    """
+
+    name: str
+    # Use current time if not provided. We need to use default factory otherwise
+    # the default value will be fixed to the build time of the class.
+    timestamp: int = field(default_factory=lambda: int(time.time() * 1e6))
+    attributes: dict[str, AttributeValue] = field(default_factory=dict)
+
+    @classmethod
+    def from_exception(cls, exception: Exception):
+        "Create a span event from an exception."
+
+        stack_trace = cls._get_stacktrace(exception)
+        return cls(
+            name="exception",
+            attributes={
+                "exception.message": str(exception),
+                "exception.type": exception.__class__.__name__,
+                "exception.stacktrace": stack_trace,
+            },
+        )
+
+    @staticmethod
+    def _get_stacktrace(error: BaseException) -> str:
+        """Get the stacktrace of the parent error."""
+        msg = repr(error)
+        try:
+            if sys.version_info < (3, 10):
+                tb = traceback.format_exception(error.__class__, error, error.__traceback__)
+            else:
+                tb = traceback.format_exception(error)
+            return "".join(tb).strip()
+        except Exception:
+            return msg
+
+    def json(self):
+        return {
+            "name": self.name,
+            "timestamp": self.timestamp,
+            "attributes": json.dumps(self.attributes, cls=CustomEncoder)
+            if self.attributes
+            else None,
+        }
+
+    def to_proto(self):
+        """Convert into OTLP compatible proto object to sent to the Databricks Trace Server."""
+        proto = ProtoSpan.Event(
+            name=self.name,
+            time_unix_nano=self.timestamp,
+        )
+        # Trace server's proto uses map<string, google.protobuf.Value> for attributes
+        for key, value in self.attributes.items():
+            set_pb_value(proto.attributes[key], value)
+        return proto
+
+
+class CustomEncoder(json.JSONEncoder):
+    """
+    Custom encoder to handle json serialization.
+    """
+
+    def default(self, o):
+        try:
+            return super().default(o)
+        except TypeError:
+            # convert datetime to string format by default
+            if isinstance(o, datetime):
+                return o.isoformat()
+            # convert object direct to string to avoid error in serialization
+            return str(o)

venv/lib/python3.9/site-packages/mlflow/entities/span_status.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from opentelemetry import trace as trace_api
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+
+class SpanStatusCode(str, Enum):
+    """Enum for status code of a span"""
+
+    # Uses the same set of status codes as OpenTelemetry
+    UNSET = "UNSET"
+    OK = "OK"
+    ERROR = "ERROR"
+
+
+@dataclass
+class SpanStatus:
+    """
+    Status of the span or the trace.
+
+    Args:
+        status_code: The status code of the span or the trace. This must be one of the
+            values of the :py:class:`mlflow.entities.SpanStatusCode` enum or a string
+            representation of it like "OK", "ERROR".
+        description: Description of the status. This should be only set when the status
+            is ERROR, otherwise it will be ignored.
+    """
+
+    status_code: SpanStatusCode
+    description: str = ""
+
+    def __post_init__(self):
+        """
+        If user provides a string status code, validate it and convert to
+        the corresponding enum value.
+        """
+        if isinstance(self.status_code, str):
+            try:
+                self.status_code = SpanStatusCode(self.status_code)
+            except ValueError:
+                raise MlflowException(
+                    f"{self.status_code} is not a valid SpanStatusCode value. "
+                    f"Please use one of {[status_code.value for status_code in SpanStatusCode]}",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+    def to_otel_status(self) -> trace_api.Status:
+        """
+        Convert :py:class:`mlflow.entities.SpanStatus` object to OpenTelemetry status object.
+
+        :meta private:
+        """
+        try:
+            status_code = getattr(trace_api.StatusCode, self.status_code.name)
+        except AttributeError:
+            raise MlflowException(
+                f"Invalid status code: {self.status_code}", error_code=INVALID_PARAMETER_VALUE
+            )
+        return trace_api.Status(status_code, self.description)
+
+    @classmethod
+    def from_otel_status(cls, otel_status: trace_api.Status) -> SpanStatus:
+        """
+        Convert OpenTelemetry status object to our status object.
+
+        :meta private:
+        """
+        try:
+            status_code = SpanStatusCode(otel_status.status_code.name)
+        except ValueError:
+            raise MlflowException(
+                f"Got invalid status code from OpenTelemetry: {otel_status.status_code}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        return cls(status_code, otel_status.description or "")

venv/lib/python3.9/site-packages/mlflow/entities/trace.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import json
+import logging
+import re
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.span import Span, SpanType
+from mlflow.entities.trace_data import TraceData
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.databricks_trace_server_pb2 import Trace as ProtoTrace
+from mlflow.protos.databricks_trace_server_pb2 import TraceData as ProtoTraceData
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Trace(_MlflowObject):
+    """A trace object.
+
+    Args:
+        info: A lightweight object that contains the metadata of a trace.
+        data: A container object that holds the spans data of a trace.
+    """
+
+    info: TraceInfo
+    data: TraceData
+
+    def __repr__(self) -> str:
+        return f"Trace(request_id={self.info.request_id})"
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"info": self.info.to_dict(), "data": self.data.to_dict()}
+
+    def to_json(self, pretty=False) -> str:
+        from mlflow.tracing.utils import TraceJSONEncoder
+
+        return json.dumps(self.to_dict(), cls=TraceJSONEncoder, indent=2 if pretty else None)
+
+    @classmethod
+    def from_dict(cls, trace_dict: dict[str, Any]) -> Trace:
+        info = trace_dict.get("info")
+        data = trace_dict.get("data")
+        if info is None or data is None:
+            raise MlflowException(
+                "Unable to parse Trace from dictionary. Expected keys: 'info' and 'data'. "
+                f"Received keys: {list(trace_dict.keys())}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        return cls(
+            info=TraceInfo.from_dict(info),
+            data=TraceData.from_dict(data),
+        )
+
+    @classmethod
+    def from_json(cls, trace_json: str) -> Trace:
+        try:
+            trace_dict = json.loads(trace_json)
+        except json.JSONDecodeError as e:
+            raise MlflowException(
+                f"Unable to parse trace JSON: {trace_json}. Error: {e}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        return cls.from_dict(trace_dict)
+
+    def _serialize_for_mimebundle(self):
+        # databricks notebooks will use the request ID to
+        # fetch the trace from the backend. including the
+        # full JSON can cause notebooks to exceed size limits
+        return json.dumps(self.info.request_id)
+
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        """
+        This method is used to trigger custom display logic in IPython notebooks.
+        See https://ipython.readthedocs.io/en/stable/config/integrating.html#MyObject
+        for more details.
+
+        At the moment, the only supported MIME type is "application/databricks.mlflow.trace",
+        which contains a JSON representation of the Trace object. This object is deserialized
+        in Databricks notebooks to display the Trace object in a nicer UI.
+        """
+        from mlflow.tracing.display import (
+            get_display_handler,
+            get_notebook_iframe_html,
+            is_using_tracking_server,
+        )
+        from mlflow.utils.databricks_utils import is_in_databricks_runtime
+
+        bundle = {"text/plain": repr(self)}
+
+        if not get_display_handler().disabled:
+            if is_in_databricks_runtime():
+                bundle["application/databricks.mlflow.trace"] = self._serialize_for_mimebundle()
+            elif is_using_tracking_server():
+                bundle["text/html"] = get_notebook_iframe_html([self])
+
+        return bundle
+
+    def to_pandas_dataframe_row(self) -> dict[str, Any]:
+        return {
+            "request_id": self.info.request_id,
+            "trace": self,
+            "timestamp_ms": self.info.timestamp_ms,
+            "status": self.info.status,
+            "execution_time_ms": self.info.execution_time_ms,
+            "request": self._deserialize_json_attr(self.data.request),
+            "response": self._deserialize_json_attr(self.data.response),
+            "request_metadata": self.info.request_metadata,
+            "spans": [span.to_dict() for span in self.data.spans],
+            "tags": self.info.tags,
+            "assessments": self.info.assessments,
+        }
+
+    def _deserialize_json_attr(self, value: str):
+        try:
+            return json.loads(value)
+        except Exception:
+            _logger.debug(f"Failed to deserialize JSON attribute: {value}", exc_info=True)
+            return value
+
+    def search_spans(
+        self, span_type: Optional[SpanType] = None, name: Optional[Union[str, re.Pattern]] = None
+    ) -> list[Span]:
+        """
+        Search for spans that match the given criteria within the trace.
+
+        Args:
+            span_type: The type of the span to search for.
+            name: The name of the span to search for. This can be a string or a regular expression.
+
+        Returns:
+            A list of spans that match the given criteria.
+            If there is no match, an empty list is returned.
+
+        .. code-block:: python
+
+            import mlflow
+            import re
+            from mlflow.entities import SpanType
+
+
+            @mlflow.trace(span_type=SpanType.CHAIN)
+            def run(x: int) -> int:
+                x = add_one(x)
+                x = add_two(x)
+                x = multiply_by_two(x)
+                return x
+
+
+            @mlflow.trace(span_type=SpanType.TOOL)
+            def add_one(x: int) -> int:
+                return x + 1
+
+
+            @mlflow.trace(span_type=SpanType.TOOL)
+            def add_two(x: int) -> int:
+                return x + 2
+
+
+            @mlflow.trace(span_type=SpanType.TOOL)
+            def multiply_by_two(x: int) -> int:
+                return x * 2
+
+
+            # Run the function and get the trace
+            y = run(2)
+            trace_id = mlflow.get_last_active_trace_id()
+            trace = mlflow.get_trace(trace_id)
+
+            # 1. Search spans by name (exact match)
+            spans = trace.search_spans(name="add_one")
+            print(spans)
+            # Output: [Span(name='add_one', ...)]
+
+            # 2. Search spans by name (regular expression)
+            pattern = re.compile(r"add.*")
+            spans = trace.search_spans(name=pattern)
+            print(spans)
+            # Output: [Span(name='add_one', ...), Span(name='add_two', ...)]
+
+            # 3. Search spans by type
+            spans = trace.search_spans(span_type=SpanType.LLM)
+            print(spans)
+            # Output: [Span(name='run', ...)]
+
+            # 4. Search spans by name and type
+            spans = trace.search_spans(name="add_one", span_type=SpanType.TOOL)
+            print(spans)
+            # Output: [Span(name='add_one', ...)]
+        """
+
+        def _match_name(span: Span) -> bool:
+            if isinstance(name, str):
+                return span.name == name
+            elif isinstance(name, re.Pattern):
+                return name.search(span.name) is not None
+            elif name is None:
+                return True
+            else:
+                raise MlflowException(
+                    f"Invalid type for 'name'. Expected str or re.Pattern. Got: {type(name)}",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+        def _match_type(span: Span) -> bool:
+            if isinstance(span_type, str):
+                return span.span_type == span_type
+            elif span_type is None:
+                return True
+            else:
+                raise MlflowException(
+                    "Invalid type for 'span_type'. Expected str or mlflow.entities.SpanType. "
+                    f"Got: {type(span_type)}",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+        return [span for span in self.data.spans if _match_name(span) and _match_type(span)]
+
+    @staticmethod
+    def pandas_dataframe_columns() -> list[str]:
+        return [
+            "request_id",
+            "trace",
+            "timestamp_ms",
+            "status",
+            "execution_time_ms",
+            "request",
+            "response",
+            "request_metadata",
+            "spans",
+            "tags",
+            "assessments",
+        ]
+
+    def to_proto(self):
+        """Convert into a proto object to sent to the Databricks Trace Server."""
+        return ProtoTrace(
+            # Convert MLflow's TraceInfoV3 to Databricks Trace Server's TraceInfo
+            info=self.info.to_v3_proto(self.data.request, self.data.response),
+            data=ProtoTraceData(spans=[span.to_proto() for span in self.data.spans]),
+        )

venv/lib/python3.9/site-packages/mlflow/entities/trace_data.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from mlflow.entities import Span
+from mlflow.tracing.constant import SpanAttributeKey
+
+
+@dataclass
+class TraceData:
+    """A container object that holds the spans data of a trace.
+
+    Args:
+        spans: List of spans that are part of the trace.
+        request: Input data for the entire trace. Equivalent to the input of the root span
+            but added for ease of access. Stored as a JSON string.
+        response: Output data for the entire trace. Equivalent to the output of the root span.
+            Stored as a JSON string.
+    """
+
+    spans: list[Span] = field(default_factory=list)
+    request: Optional[str] = None
+    response: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, d):
+        if not isinstance(d, dict):
+            raise TypeError(f"TraceData.from_dict() expects a dictionary. Got: {type(d).__name__}")
+        return cls(
+            request=d.get("request"),
+            response=d.get("response"),
+            spans=[Span.from_dict(span) for span in d.get("spans", [])],
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "spans": [span.to_dict() for span in self.spans],
+            "request": self.request,
+            "response": self.response,
+        }
+
+    @property
+    def intermediate_outputs(self) -> Optional[dict[str, Any]]:
+        """
+        Returns intermediate outputs produced by the model or agent while handling the request.
+        There are mainly two flows to return intermediate outputs:
+        1. When a trace is generate by the `mlflow.log_trace` API,
+        return `intermediate_outputs` attribute of the span.
+        2. When a trace is created normally with a tree of spans,
+        aggregate the outputs of non-root spans.
+        """
+        root_span = self._get_root_span()
+        if root_span and root_span.get_attribute(SpanAttributeKey.INTERMEDIATE_OUTPUTS):
+            return root_span.get_attribute(SpanAttributeKey.INTERMEDIATE_OUTPUTS)
+
+        if len(self.spans) > 1:
+            return {
+                span.name: span.outputs
+                for span in self.spans
+                if span.parent_id and span.outputs is not None
+            }
+
+    def _get_root_span(self) -> Optional[Span]:
+        for span in self.spans:
+            if span.parent_id is None:
+                return span

venv/lib/python3.9/site-packages/mlflow/entities/trace_info.py

@@ -0,0 +1,148 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any, Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.assessment import Assessment
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.protos.databricks_trace_server_pb2 import TraceInfo as ProtoTraceInfoV3
+from mlflow.protos.service_pb2 import TraceInfo as ProtoTraceInfo
+from mlflow.protos.service_pb2 import TraceLocation as ProtoTraceLocation
+from mlflow.protos.service_pb2 import TraceRequestMetadata as ProtoTraceRequestMetadata
+from mlflow.protos.service_pb2 import TraceTag as ProtoTraceTag
+
+
+def _truncate_request_metadata(d: dict[str, Any]) -> dict[str, str]:
+    from mlflow.tracing.constant import MAX_CHARS_IN_TRACE_INFO_METADATA
+
+    return {
+        k[:MAX_CHARS_IN_TRACE_INFO_METADATA]: str(v)[:MAX_CHARS_IN_TRACE_INFO_METADATA]
+        for k, v in d.items()
+    }
+
+
+def _truncate_tags(d: dict[str, Any]) -> dict[str, str]:
+    from mlflow.tracing.constant import (
+        MAX_CHARS_IN_TRACE_INFO_TAGS_KEY,
+        MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE,
+    )
+
+    return {
+        k[:MAX_CHARS_IN_TRACE_INFO_TAGS_KEY]: str(v)[:MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE]
+        for k, v in d.items()
+    }
+
+
+@dataclass
+class TraceInfo(_MlflowObject):
+    """Metadata about a trace.
+
+    Args:
+        request_id: id of the trace.
+        experiment_id: id of the experiment.
+        timestamp_ms: start time of the trace, in milliseconds.
+        execution_time_ms: duration of the trace, in milliseconds.
+        status: status of the trace.
+        request_metadata: Key-value pairs associated with the trace. Request metadata are designed
+            for immutable values like run ID associated with the trace.
+        tags: Tags associated with the trace. Tags are designed for mutable values like trace name,
+            that can be updated by the users after the trace is created, unlike request_metadata.
+    """
+
+    request_id: str
+    experiment_id: str
+    timestamp_ms: int
+    execution_time_ms: Optional[int]
+    status: TraceStatus
+    request_metadata: dict[str, str] = field(default_factory=dict)
+    tags: dict[str, str] = field(default_factory=dict)
+    assessments: list[Assessment] = field(default_factory=list)
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    def to_proto(self):
+        proto = ProtoTraceInfo()
+        proto.request_id = self.request_id
+        proto.experiment_id = self.experiment_id
+        proto.timestamp_ms = self.timestamp_ms
+        # NB: Proto setter does not support nullable fields (even with 'optional' keyword),
+        # so we substitute None with 0 for execution_time_ms. This should be not too confusing
+        # as we only put None when starting a trace i.e. the execution time is actually 0.
+        proto.execution_time_ms = self.execution_time_ms or 0
+        proto.status = self.status.to_proto()
+
+        request_metadata = []
+        for key, value in _truncate_request_metadata(self.request_metadata).items():
+            attr = ProtoTraceRequestMetadata()
+            attr.key = key
+            attr.value = value
+            request_metadata.append(attr)
+        proto.request_metadata.extend(request_metadata)
+
+        tags = []
+        for key, value in _truncate_tags(self.tags).items():
+            tag = ProtoTraceTag()
+            tag.key = key
+            tag.value = str(value)
+            tags.append(tag)
+
+        proto.tags.extend(tags)
+        return proto
+
+    @classmethod
+    def from_proto(cls, proto, assessments=None):
+        return cls(
+            request_id=proto.request_id,
+            experiment_id=proto.experiment_id,
+            timestamp_ms=proto.timestamp_ms,
+            execution_time_ms=proto.execution_time_ms,
+            status=TraceStatus.from_proto(proto.status),
+            request_metadata={attr.key: attr.value for attr in proto.request_metadata},
+            tags={tag.key: tag.value for tag in proto.tags},
+            assessments=assessments or [],
+        )
+
+    def to_dict(self):
+        """
+        Convert trace info to a dictionary for persistence.
+        Update status field to the string value for serialization.
+        """
+        trace_info_dict = asdict(self)
+        trace_info_dict["status"] = self.status.value
+        return trace_info_dict
+
+    @classmethod
+    def from_dict(cls, trace_info_dict):
+        """
+        Convert trace info dictionary to TraceInfo object.
+        """
+        if "status" not in trace_info_dict:
+            raise ValueError("status is required in trace info dictionary.")
+        trace_info_dict["status"] = TraceStatus(trace_info_dict["status"])
+        return cls(**trace_info_dict)
+
+    def to_v3_proto(self, request: Optional[str], response: Optional[str]):
+        """Convert into the V3 TraceInfo proto object."""
+        proto = ProtoTraceInfoV3()
+
+        proto.trace_id = self.request_id
+        proto.trace_location.type = ProtoTraceLocation.MLFLOW_EXPERIMENT
+        proto.trace_location.mlflow_experiment.experiment_id = self.experiment_id
+
+        proto.request = request or ""
+        proto.response = response or ""
+        proto.state = ProtoTraceInfoV3.State.Value(self.status.name)
+
+        proto.request_time.FromMilliseconds(self.timestamp_ms)
+        if self.execution_time_ms is not None:
+            proto.execution_duration.FromMilliseconds(self.execution_time_ms)
+
+        if self.request_metadata:
+            proto.trace_metadata.update(_truncate_request_metadata(self.request_metadata))
+
+        if self.tags:
+            proto.tags.update(_truncate_tags(self.tags))
+
+        return proto

venv/lib/python3.9/site-packages/mlflow/entities/trace_status.py

@@ -0,0 +1,42 @@
+from enum import Enum
+
+from opentelemetry import trace as trace_api
+
+from mlflow.protos.service_pb2 import TraceStatus as ProtoTraceStatus
+
+
+class TraceStatus(str, Enum):
+    """Enum for status of an :py:class:`mlflow.entities.TraceInfo`."""
+
+    UNSPECIFIED = "TRACE_STATUS_UNSPECIFIED"
+    OK = "OK"
+    ERROR = "ERROR"
+    IN_PROGRESS = "IN_PROGRESS"
+
+    def to_proto(self):
+        return ProtoTraceStatus.Value(self)
+
+    @staticmethod
+    def from_proto(proto_status):
+        return TraceStatus(ProtoTraceStatus.Name(proto_status))
+
+    @staticmethod
+    def from_otel_status(otel_status: trace_api.Status):
+        return _OTEL_STATUS_CODE_TO_MLFLOW[otel_status.status_code]
+
+    @classmethod
+    def pending_statuses(cls):
+        """Traces in pending statuses can be updated to any statuses."""
+        return {cls.IN_PROGRESS}
+
+    @classmethod
+    def end_statuses(cls):
+        """Traces in end statuses cannot be updated to any statuses."""
+        return {cls.UNSPECIFIED, cls.OK, cls.ERROR}
+
+
+_OTEL_STATUS_CODE_TO_MLFLOW = {
+    trace_api.StatusCode.OK: TraceStatus.OK,
+    trace_api.StatusCode.ERROR: TraceStatus.ERROR,
+    trace_api.StatusCode.UNSET: TraceStatus.UNSPECIFIED,
+}

venv/lib/python3.9/site-packages/mlflow/entities/view_type.py

@@ -0,0 +1,51 @@
+from mlflow.protos import service_pb2
+
+
+class ViewType:
+    """Enum to filter requested experiment types."""
+
+    ACTIVE_ONLY, DELETED_ONLY, ALL = range(1, 4)
+    _VIEW_TO_STRING = {
+        ACTIVE_ONLY: "active_only",
+        DELETED_ONLY: "deleted_only",
+        ALL: "all",
+    }
+    _STRING_TO_VIEW = {value: key for key, value in _VIEW_TO_STRING.items()}
+
+    @classmethod
+    def from_string(cls, view_str):
+        if view_str not in cls._STRING_TO_VIEW:
+            raise Exception(
+                f"Could not get valid view type corresponding to string {view_str}. "
+                f"Valid view types are {list(cls._STRING_TO_VIEW.keys())}"
+            )
+        return cls._STRING_TO_VIEW[view_str]
+
+    @classmethod
+    def to_string(cls, view_type):
+        if view_type not in cls._VIEW_TO_STRING:
+            raise Exception(
+                f"Could not get valid view type corresponding to string {view_type}. "
+                f"Valid view types are {list(cls._VIEW_TO_STRING.keys())}"
+            )
+        return cls._VIEW_TO_STRING[view_type]
+
+    @classmethod
+    def to_proto(cls, view_type):
+        if view_type == cls.ACTIVE_ONLY:
+            return service_pb2.ACTIVE_ONLY
+        elif view_type == cls.DELETED_ONLY:
+            return service_pb2.DELETED_ONLY
+        elif view_type == cls.ALL:
+            return service_pb2.ALL
+        raise ValueError(f"Unexpected view_type: {view_type}")
+
+    @classmethod
+    def from_proto(cls, proto_view_type):
+        if proto_view_type == service_pb2.ACTIVE_ONLY:
+            return cls.ACTIVE_ONLY
+        elif proto_view_type == service_pb2.DELETED_ONLY:
+            return cls.DELETED_ONLY
+        elif proto_view_type == service_pb2.ALL:
+            return cls.ALL
+        raise ValueError(f"Unexpected proto_view_type: {proto_view_type}")