genesis-flow 1.0.0__py3-none-any.whl

mlflow/entities/trace_info.py

@@ -0,0 +1,220 @@
+import json
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from google.protobuf.duration_pb2 import Duration
+from google.protobuf.json_format import MessageToDict
+from google.protobuf.timestamp_pb2 import Timestamp
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.entities.assessment import Assessment
+from mlflow.entities.trace_location import TraceLocation
+from mlflow.entities.trace_state import TraceState
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.protos.service_pb2 import TraceInfoV3 as ProtoTraceInfoV3
+from mlflow.tracing.constant import TRACE_SCHEMA_VERSION, TRACE_SCHEMA_VERSION_KEY, TraceMetadataKey
+
+
+@dataclass
+class TraceInfo(_MlflowObject):
+    """Metadata about a trace, such as its ID, location, timestamp, etc.
+
+    Args:
+        trace_id: The primary identifier for the trace.
+        trace_location: The location where the trace is stored, represented as
+            a :py:class:`~mlflow.entities.TraceLocation` object. MLflow currently
+            support MLflow Experiment or Databricks Inference Table as a trace location.
+        request_time: Start time of the trace, in milliseconds.
+        state: State of the trace, represented as a :py:class:`~mlflow.entities.TraceState`
+            enum. Can be one of [`OK`, `ERROR`, `IN_PROGRESS`, `STATE_UNSPECIFIED`].
+        request_preview: Request to the model/agent, equivalent to the input of the root,
+            span but JSON-encoded and can be truncated.
+        response_preview: Response from the model/agent, equivalent to the output of the
+            root span but JSON-encoded and can be truncated.
+        client_request_id: Client supplied request ID associated with the trace. This
+            could be used to identify the trace/request from an external system that
+            produced the trace, e.g., a session ID in a web application.
+        execution_duration: Duration of the trace, in milliseconds.
+        trace_metadata: Key-value pairs associated with the trace. They are designed
+            for immutable values like run ID associated with the trace.
+        tags: Tags associated with the trace. They are designed for mutable values,
+            that can be updated after the trace is created via MLflow UI or API.
+        assessments: List of assessments associated with the trace.
+    """
+
+    trace_id: str
+    trace_location: TraceLocation
+    request_time: int
+    state: TraceState
+    request_preview: Optional[str] = None
+    response_preview: Optional[str] = None
+    client_request_id: Optional[str] = None
+    execution_duration: Optional[int] = None
+    trace_metadata: dict[str, str] = field(default_factory=dict)
+    tags: dict[str, str] = field(default_factory=dict)
+    assessments: list[Assessment] = field(default_factory=list)
+
+    def __post_init__(self):
+        # NB: MLflow automatically converts trace metadata and spans to V3 format, even if the
+        # trace was originally created in V2 format with an earlier version of MLflow. Accordingly,
+        # we also update the `TRACE_SCHEMA_VERSION_KEY` in the trace metadata to V3 for consistency
+        self.trace_metadata[TRACE_SCHEMA_VERSION_KEY] = str(TRACE_SCHEMA_VERSION)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the TraceInfoV3 object to a dictionary."""
+        res = MessageToDict(self.to_proto(), preserving_proto_field_name=True)
+        if self.execution_duration is not None:
+            res.pop("execution_duration", None)
+            res["execution_duration_ms"] = self.execution_duration
+        return res
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "TraceInfo":
+        """Create a TraceInfoV3 object from a dictionary."""
+        if "request_id" in d:
+            from mlflow.entities.trace_info_v2 import TraceInfoV2
+
+            return TraceInfoV2.from_dict(d).to_v3()
+
+        d = d.copy()
+        if assessments := d.get("assessments"):
+            d["assessments"] = [Assessment.from_dictionary(a) for a in assessments]
+
+        if trace_location := d.get("trace_location"):
+            d["trace_location"] = TraceLocation.from_dict(trace_location)
+
+        if state := d.get("state"):
+            d["state"] = TraceState(state)
+
+        if request_time := d.get("request_time"):
+            timestamp = Timestamp()
+            timestamp.FromJsonString(request_time)
+            d["request_time"] = timestamp.ToMilliseconds()
+
+        if (execution_duration := d.pop("execution_duration_ms", None)) is not None:
+            d["execution_duration"] = execution_duration
+
+        return cls(**d)
+
+    def to_proto(self):
+        from mlflow.entities.trace_info_v2 import _truncate_request_metadata, _truncate_tags
+
+        request_time = Timestamp()
+        request_time.FromMilliseconds(self.request_time)
+        execution_duration = None
+        if self.execution_duration is not None:
+            execution_duration = Duration()
+            execution_duration.FromMilliseconds(self.execution_duration)
+
+        return ProtoTraceInfoV3(
+            trace_id=self.trace_id,
+            client_request_id=self.client_request_id,
+            trace_location=self.trace_location.to_proto(),
+            request_preview=self.request_preview,
+            response_preview=self.response_preview,
+            request_time=request_time,
+            execution_duration=execution_duration,
+            state=self.state.to_proto(),
+            trace_metadata=_truncate_request_metadata(self.trace_metadata),
+            tags=_truncate_tags(self.tags),
+            assessments=[a.to_proto() for a in self.assessments],
+        )
+
+    @classmethod
+    def from_proto(cls, proto) -> "TraceInfo":
+        if "request_id" in proto.DESCRIPTOR.fields_by_name:
+            from mlflow.entities.trace_info_v2 import TraceInfoV2
+
+            return TraceInfoV2.from_proto(proto).to_v3()
+
+        return cls(
+            trace_id=proto.trace_id,
+            client_request_id=(
+                proto.client_request_id if proto.HasField("client_request_id") else None
+            ),
+            trace_location=TraceLocation.from_proto(proto.trace_location),
+            request_preview=proto.request_preview,
+            response_preview=proto.response_preview,
+            request_time=proto.request_time.ToMilliseconds(),
+            execution_duration=(
+                proto.execution_duration.ToMilliseconds()
+                if proto.HasField("execution_duration")
+                else None
+            ),
+            state=TraceState.from_proto(proto.state),
+            trace_metadata=dict(proto.trace_metadata),
+            tags=dict(proto.tags),
+            assessments=[Assessment.from_proto(a) for a in proto.assessments],
+        )
+
+    # Aliases for backward compatibility with V2 format
+    @property
+    def request_id(self) -> str:
+        """Deprecated. Use `trace_id` instead."""
+        return self.trace_id
+
+    @property
+    def experiment_id(self) -> Optional[str]:
+        """
+        An MLflow experiment ID associated with the trace, if the trace is stored
+        in MLflow tracking server. Otherwise, None.
+        """
+        return (
+            self.trace_location.mlflow_experiment
+            and self.trace_location.mlflow_experiment.experiment_id
+        )
+
+    @experiment_id.setter
+    def experiment_id(self, value: Optional[str]) -> None:
+        self.trace_location.mlflow_experiment.experiment_id = value
+
+    @property
+    def request_metadata(self) -> dict[str, str]:
+        """Deprecated. Use `trace_metadata` instead."""
+        return self.trace_metadata
+
+    @property
+    def timestamp_ms(self) -> int:
+        return self.request_time
+
+    @timestamp_ms.setter
+    def timestamp_ms(self, value: int) -> None:
+        self.request_time = value
+
+    @property
+    def execution_time_ms(self) -> Optional[int]:
+        return self.execution_duration
+
+    @execution_time_ms.setter
+    def execution_time_ms(self, value: Optional[int]) -> None:
+        self.execution_duration = value
+
+    @property
+    def status(self) -> TraceStatus:
+        """Deprecated. Use `state` instead."""
+        return TraceStatus.from_state(self.state)
+
+    @status.setter
+    def status(self, value: TraceStatus) -> None:
+        self.state = value.to_state()
+
+    @property
+    def token_usage(self) -> Optional[dict[str, int]]:
+        """
+        Returns the aggregated token usage for the trace.
+
+        Returns:
+            A dictionary containing the aggregated LLM token usage for the trace.
+            - "input_tokens": The total number of input tokens.
+            - "output_tokens": The total number of output tokens.
+            - "total_tokens": Sum of input and output tokens.
+
+        .. note::
+
+            The token usage tracking is not supported for all LLM providers.
+            Refer to the MLflow Tracing documentation for which providers
+            support token usage tracking.
+        """
+        if usage_json := self.trace_metadata.get(TraceMetadataKey.TOKEN_USAGE):
+            return json.loads(usage_json)
+        return None

mlflow/entities/trace_info_v2.py

@@ -0,0 +1,162 @@
+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_info import TraceInfo
+from mlflow.entities.trace_location import TraceLocation
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.protos.service_pb2 import TraceInfo as ProtoTraceInfo
+from mlflow.protos.service_pb2 import TraceRequestMetadata as ProtoTraceRequestMetadata
+from mlflow.protos.service_pb2 import TraceTag as ProtoTraceTag
+from mlflow.tracing.constant import TRACE_SCHEMA_VERSION_KEY
+
+
+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 TraceInfoV2(_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 __post_init__(self):
+        self.request_metadata[TRACE_SCHEMA_VERSION_KEY] = "2"
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def trace_id(self) -> str:
+        """Returns the trace ID of the trace info."""
+        return self.request_id
+
+    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
+        # Client request ID field is only added for internal use, and should not be
+        # serialized for V2 TraceInfo.
+        trace_info_dict.pop("client_request_id", None)
+        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(self, request: Optional[str] = None, response: Optional[str] = None) -> TraceInfo:
+        return TraceInfo(
+            trace_id=self.request_id,
+            trace_location=TraceLocation.from_experiment_id(self.experiment_id),
+            request_preview=request,
+            response_preview=response,
+            request_time=self.timestamp_ms,
+            execution_duration=self.execution_time_ms,
+            state=self.status.to_state(),
+            trace_metadata=self.request_metadata.copy(),
+            tags=self.tags,
+            assessments=self.assessments,
+        )
+
+    @classmethod
+    def from_v3(cls, trace_info: TraceInfo) -> "TraceInfoV2":
+        return cls(
+            request_id=trace_info.trace_id,
+            experiment_id=trace_info.experiment_id,
+            timestamp_ms=trace_info.request_time,
+            execution_time_ms=trace_info.execution_duration,
+            status=TraceStatus.from_state(trace_info.state),
+            request_metadata=trace_info.trace_metadata.copy(),
+            tags=trace_info.tags,
+        )

mlflow/entities/trace_location.py

@@ -0,0 +1,173 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.exceptions import MlflowException
+from mlflow.protos import service_pb2 as pb
+
+
+@dataclass
+class MlflowExperimentLocation(_MlflowObject):
+    """
+    Represents the location of an MLflow experiment.
+
+    Args:
+        experiment_id: The ID of the MLflow experiment where the trace is stored.
+    """
+
+    experiment_id: str
+
+    def to_proto(self):
+        return pb.TraceLocation.MlflowExperimentLocation(experiment_id=self.experiment_id)
+
+    @classmethod
+    def from_proto(cls, proto) -> "MlflowExperimentLocation":
+        return cls(experiment_id=proto.experiment_id)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"experiment_id": self.experiment_id}
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "MlflowExperimentLocation":
+        return cls(experiment_id=d["experiment_id"])
+
+
+@dataclass
+class InferenceTableLocation(_MlflowObject):
+    """
+    Represents the location of a Databricks inference table.
+
+    Args:
+        full_table_name: The fully qualified name of the inference table where
+            the trace is stored, in the format of `<catalog>.<schema>.<table>`.
+    """
+
+    full_table_name: str
+
+    def to_proto(self):
+        return pb.TraceLocation.InferenceTableLocation(full_table_name=self.full_table_name)
+
+    @classmethod
+    def from_proto(cls, proto) -> "InferenceTableLocation":
+        return cls(full_table_name=proto.full_table_name)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"full_table_name": self.full_table_name}
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "InferenceTableLocation":
+        return cls(full_table_name=d["full_table_name"])
+
+
+class TraceLocationType(str, Enum):
+    TRACE_LOCATION_TYPE_UNSPECIFIED = "TRACE_LOCATION_TYPE_UNSPECIFIED"
+    MLFLOW_EXPERIMENT = "MLFLOW_EXPERIMENT"
+    INFERENCE_TABLE = "INFERENCE_TABLE"
+
+    def to_proto(self):
+        return pb.TraceLocation.TraceLocationType.Value(self)
+
+    @classmethod
+    def from_proto(cls, proto: int) -> "TraceLocationType":
+        return TraceLocationType(pb.TraceLocation.TraceLocationType.Name(proto))
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "TraceLocationType":
+        return cls(d["type"])
+
+
+@dataclass
+class TraceLocation(_MlflowObject):
+    """
+    Represents the location where the trace is stored.
+
+    Currently, MLflow supports two types of trace locations:
+
+        - MLflow experiment: The trace is stored in an MLflow experiment.
+        - Inference table: The trace is stored in a Databricks inference table.
+
+    Args:
+        type: The type of the trace location, should be one of the
+            :py:class:`TraceLocationType` enum values.
+        mlflow_experiment: The MLflow experiment location. Set this when the
+            location type is MLflow experiment.
+        inference_table: The inference table location. Set this when the
+            location type is Databricks Inference table.
+    """
+
+    type: TraceLocationType
+    mlflow_experiment: Optional[MlflowExperimentLocation] = None
+    inference_table: Optional[InferenceTableLocation] = None
+
+    def __post_init__(self) -> None:
+        if self.mlflow_experiment is not None and self.inference_table is not None:
+            raise MlflowException.invalid_parameter_value(
+                "Only one of mlflow_experiment or inference_table can be provided."
+            )
+
+        if (self.mlflow_experiment and self.type != TraceLocationType.MLFLOW_EXPERIMENT) or (
+            self.inference_table and self.type != TraceLocationType.INFERENCE_TABLE
+        ):
+            raise MlflowException.invalid_parameter_value(
+                f"Trace location type {type} does not match the provided location "
+                f"{self.mlflow_experiment or self.inference_table}."
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        d = {"type": self.type.value}
+        if self.mlflow_experiment:
+            d["mlflow_experiment"] = self.mlflow_experiment.to_dict()
+        elif self.inference_table:
+            d["inference_table"] = self.inference_table.to_dict()
+        return d
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "TraceLocation":
+        return cls(
+            type=TraceLocationType(d["type"]),
+            mlflow_experiment=(
+                MlflowExperimentLocation.from_dict(v) if (v := d.get("mlflow_experiment")) else None
+            ),
+            inference_table=(
+                InferenceTableLocation.from_dict(v) if (v := d.get("inference_table")) else None
+            ),
+        )
+
+    def to_proto(self):
+        if self.mlflow_experiment:
+            return pb.TraceLocation(
+                type=self.type.to_proto(),
+                mlflow_experiment=self.mlflow_experiment.to_proto(),
+            )
+        elif self.inference_table:
+            return pb.TraceLocation(
+                type=self.type.to_proto(),
+                inference_table=self.inference_table.to_proto(),
+            )
+
+        else:
+            return pb.TraceLocation(type=self.type.to_proto())
+
+    @classmethod
+    def from_proto(cls, proto) -> "TraceLocation":
+        type_ = TraceLocationType.from_proto(proto.type)
+        if proto.WhichOneof("identifier") == "mlflow_experiment":
+            return cls(
+                type=type_,
+                mlflow_experiment=MlflowExperimentLocation.from_proto(proto.mlflow_experiment),
+            )
+        elif proto.WhichOneof("identifier") == "inference_table":
+            return cls(
+                type=type_,
+                inference_table=InferenceTableLocation.from_proto(proto.inference_table),
+            )
+        else:
+            return cls(type=type_)
+
+    @classmethod
+    def from_experiment_id(cls, experiment_id: str) -> "TraceLocation":
+        return cls(
+            type=TraceLocationType.MLFLOW_EXPERIMENT,
+            mlflow_experiment=MlflowExperimentLocation(experiment_id=experiment_id),
+        )

mlflow/entities/trace_state.py

@@ -0,0 +1,39 @@
+from enum import Enum
+
+from opentelemetry import trace as trace_api
+
+from mlflow.protos import service_pb2 as pb
+
+
+class TraceState(str, Enum):
+    """Enum representing the state of a trace.
+
+    - ``STATE_UNSPECIFIED``: Unspecified trace state.
+    - ``OK``: Trace successfully completed.
+    - ``ERROR``: Trace encountered an error.
+    - ``IN_PROGRESS``: Trace is currently in progress.
+    """
+
+    STATE_UNSPECIFIED = "STATE_UNSPECIFIED"
+    OK = "OK"
+    ERROR = "ERROR"
+    IN_PROGRESS = "IN_PROGRESS"
+
+    def to_proto(self):
+        return pb.TraceInfoV3.State.Value(self)
+
+    @classmethod
+    def from_proto(cls, proto: int) -> "TraceState":
+        return TraceState(pb.TraceInfoV3.State.Name(proto))
+
+    @staticmethod
+    def from_otel_status(otel_status: trace_api.Status):
+        """Convert OpenTelemetry status code to MLflow TraceState."""
+        return _OTEL_STATUS_CODE_TO_MLFLOW[otel_status.status_code]
+
+
+_OTEL_STATUS_CODE_TO_MLFLOW = {
+    trace_api.StatusCode.OK: TraceState.OK,
+    trace_api.StatusCode.ERROR: TraceState.ERROR,
+    trace_api.StatusCode.UNSET: TraceState.STATE_UNSPECIFIED,
+}

mlflow/entities/trace_status.py

@@ -0,0 +1,68 @@
+from enum import Enum
+
+from opentelemetry import trace as trace_api
+
+from mlflow.entities.trace_state import TraceState
+from mlflow.protos.service_pb2 import TraceStatus as ProtoTraceStatus
+from mlflow.utils.annotations import deprecated
+
+
+@deprecated(alternative="mlflow.entities.trace_state.TraceState")
+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_state(self) -> TraceState:
+        if self == TraceStatus.UNSPECIFIED:
+            return TraceState.STATE_UNSPECIFIED
+        elif self == TraceStatus.OK:
+            return TraceState.OK
+        elif self == TraceStatus.ERROR:
+            return TraceState.ERROR
+        elif self == TraceStatus.IN_PROGRESS:
+            return TraceState.IN_PROGRESS
+        raise ValueError(f"Unknown TraceStatus: {self}")
+
+    @classmethod
+    def from_state(cls, state: TraceState) -> "TraceStatus":
+        if state == TraceState.STATE_UNSPECIFIED:
+            return cls.UNSPECIFIED
+        elif state == TraceState.OK:
+            return cls.OK
+        elif state == TraceState.ERROR:
+            return cls.ERROR
+        elif state == TraceState.IN_PROGRESS:
+            return cls.IN_PROGRESS
+        raise ValueError(f"Unknown TraceState: {state}")
+
+    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,
+}