genesis-flow 1.0.0__py3-none-any.whl

mlflow/entities/span_event.py

@@ -0,0 +1,96 @@
+import json
+import sys
+import time
+import traceback
+from dataclasses import dataclass, field
+from datetime import datetime
+
+from google.protobuf.json_format import ParseDict
+from google.protobuf.struct_pb2 import Value
+from opentelemetry.util.types import AttributeValue
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.databricks_trace_server_pb2 import Span as ProtoSpan
+
+
+@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 nanoseconds.
+            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."""
+        return ProtoSpan.Event(
+            name=self.name,
+            time_unix_nano=self.timestamp,
+            attributes={k: ParseDict(v, Value()) for k, v in self.attributes.items()},
+        )
+
+
+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)

mlflow/entities/span_status.py

@@ -0,0 +1,102 @@
+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"
+
+    @staticmethod
+    def from_proto_status_code(status_code: str) -> SpanStatusCode:
+        """
+        Convert a string status code to the corresponding SpanStatusCode enum value.
+        """
+        from mlflow.protos.databricks_trace_server_pb2 import Span as ProtoSpan
+
+        proto_code = ProtoSpan.Status.StatusCode
+        status_code_mapping = {
+            proto_code.Name(proto_code.STATUS_CODE_UNSET): SpanStatusCode.UNSET,
+            proto_code.Name(proto_code.STATUS_CODE_OK): SpanStatusCode.OK,
+            proto_code.Name(proto_code.STATUS_CODE_ERROR): SpanStatusCode.ERROR,
+        }
+        try:
+            return status_code_mapping[status_code]
+        except KeyError:
+            raise MlflowException(
+                f"Invalid status code: {status_code}. "
+                f"Valid values are: {', '.join(status_code_mapping.keys())}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+
+@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 "")

mlflow/entities/trace.py

@@ -0,0 +1,317 @@
+from __future__ import annotations
+
+import json
+import logging
+import re
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal, 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.entities.trace_info_v2 import TraceInfoV2
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.service_pb2 import Trace as ProtoTrace
+
+if TYPE_CHECKING:
+    from mlflow.entities.assessment import Assessment
+
+_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 __post_init__(self):
+        if isinstance(self.info, TraceInfoV2):
+            self.info = self.info.to_v3(request=self.data.request, response=self.data.response)
+
+    def __repr__(self) -> str:
+        return f"Trace(trace_id={self.info.trace_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 {
+            "trace_id": self.info.trace_id,
+            "trace": self.to_json(),  # json string to be compatible with Spark DataFrame
+            "client_request_id": self.info.client_request_id,
+            "state": self.info.state,
+            "request_time": self.info.request_time,
+            "execution_duration": self.info.execution_duration,
+            "request": self._deserialize_json_attr(self.data.request),
+            "response": self._deserialize_json_attr(self.data.response),
+            "trace_metadata": self.info.trace_metadata,
+            "tags": self.info.tags,
+            "spans": [span.to_dict() for span in self.data.spans],
+            "assessments": [assessment.to_dictionary() for assessment in 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,
+        span_id: Optional[str] = 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.
+            span_id: The ID of the span to search for.
+
+        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,
+                )
+
+        def _match_id(span: Span) -> bool:
+            if span_id is None:
+                return True
+            else:
+                return span.span_id == span_id
+
+        return [
+            span
+            for span in self.data.spans
+            if _match_name(span) and _match_type(span) and _match_id(span)
+        ]
+
+    def search_assessments(
+        self,
+        name: Optional[str] = None,
+        *,
+        span_id: Optional[str] = None,
+        all: bool = False,
+        type: Optional[Literal["expectation", "feedback"]] = None,
+    ) -> list["Assessment"]:
+        """
+        Get assessments for a given name / span ID. By default, this only returns assessments
+        that are valid (i.e. have not been overridden by another assessment). To return all
+        assessments, specify `all=True`.
+
+        Args:
+            name: The name of the assessment to get. If not provided, this will match
+                all assessment names.
+            span_id: The span ID to get assessments for.
+                If not provided, this will match all spans.
+            all: If True, return all assessments regardless of validity.
+            type: The type of assessment to get (one of "feedback" or "expectation").
+                If not provided, this will match all assessment types.
+
+        Returns:
+            A list of assessments that meet the given conditions.
+        """
+
+        def validate_type(assessment: Assessment) -> bool:
+            from mlflow.entities.assessment import Expectation, Feedback
+
+            if type == "expectation":
+                return isinstance(assessment, Expectation)
+            elif type == "feedback":
+                return isinstance(assessment, Feedback)
+
+            return True
+
+        return [
+            assessment
+            for assessment in self.info.assessments
+            if (name is None or assessment.name == name)
+            and (span_id is None or assessment.span_id == span_id)
+            # valid defaults to true, so Nones are valid
+            and (all or assessment.valid in (True, None))
+            and (type is None or validate_type(assessment))
+        ]
+
+    @staticmethod
+    def pandas_dataframe_columns() -> list[str]:
+        return [
+            "trace_id",
+            "trace",
+            "client_request_id",
+            "state",
+            "request_time",
+            "execution_duration",
+            "request",
+            "response",
+            "trace_metadata",
+            "tags",
+            "spans",
+            "assessments",
+        ]
+
+    def to_proto(self):
+        """
+        Convert into a proto object to sent to the MLflow backend.
+
+        NB: The Trace definition in MLflow backend doesn't include the `data` field,
+            but rather only contains TraceInfoV3.
+        """
+
+        return ProtoTrace(trace_info=self.info.to_proto())

mlflow/entities/trace_data.py

@@ -0,0 +1,71 @@
+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.
+    """
+
+    spans: list[Span] = field(default_factory=list)
+
+    # NB: Custom constructor to allow passing additional kwargs for backward compatibility for
+    # DBX agent evaluator. Once they migrates to trace V3 schema, we can remove this.
+    def __init__(self, spans: Optional[list[Span]] = None, **kwargs):
+        self.spans = spans or []
+
+    @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(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]}
+
+    @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
+
+    # `request` and `response` are preserved for backward compatibility with v2
+    @property
+    def request(self) -> Optional[str]:
+        if span := self._get_root_span():
+            # Accessing the OTel span directly get serialized value directly.
+            return span._span.attributes.get(SpanAttributeKey.INPUTS)
+        return None
+
+    @property
+    def response(self) -> Optional[str]:
+        if span := self._get_root_span():
+            # Accessing the OTel span directly get serialized value directly.
+            return span._span.attributes.get(SpanAttributeKey.OUTPUTS)
+        return None