mlflow--mlflow/mlflow/tracing/constant.py

from enum import Enum


# NB: These keys are placeholders and subject to change
class TraceMetadataKey:
    INPUTS = "mlflow.traceInputs"
    OUTPUTS = "mlflow.traceOutputs"
    SOURCE_RUN = "mlflow.sourceRun"
    MODEL_ID = "mlflow.modelId"
    # Trace size statistics including total size, number of spans, and max span size
    SIZE_STATS = "mlflow.trace.sizeStats"
    # Aggregated token usage information in a single trace, stored as a dumped JSON string.
    TOKEN_USAGE = "mlflow.trace.tokenUsage"
    # Set by start_trace() when it writes the authoritative (DFS-dedup) TOKEN_USAGE / COST
    # so that concurrent log_spans() calls do not accumulate on top of them.
    # Set by start_trace() after writing authoritative trace-level values (TOKEN_USAGE,
    # COST, session ID, request_time, execution_duration) so that concurrent log_spans()
    # calls do not overwrite them.
    TRACE_INFO_FINALIZED = "mlflow.trace.infoFinalized"
    # Aggregated cost information in a single trace, stored as a dumped JSON string (USD).
    COST = "mlflow.trace.cost"
    # Store the user ID/name of the application request. Do not confuse this with mlflow.user
    # tag, which stores "who created the trace" i.e. developer or system name.
    TRACE_USER = "mlflow.trace.user"
    # Store the session ID of the application request.
    TRACE_SESSION = "mlflow.trace.session"

    # Total size of the trace in bytes. Deprecated, use SIZE_STATS instead.
    SIZE_BYTES = "mlflow.trace.sizeBytes"

    # Gateway-specific metadata keys
    GATEWAY_ENDPOINT_ID = "mlflow.gateway.endpointId"
    GATEWAY_REQUEST_TYPE = "mlflow.gateway.requestType"
    GATEWAY_CALLER = "mlflow.gateway.caller"
    # Store the user ID/name from authentication
    AUTH_USER_ID = "mlflow.auth.userId"
    AUTH_USERNAME = "mlflow.auth.username"


class TraceTagKey:
    TRACE_NAME = "mlflow.traceName"
    EVAL_REQUEST_ID = "mlflow.eval.requestId"
    SPANS_LOCATION = "mlflow.trace.spansLocation"
    ARCHIVE_LOCATION = "mlflow.trace.archiveLocation"
    ARCHIVAL_FAILURE = "mlflow.trace.archivalFailure"
    # Store the source scorer name that generated the trace. This tag is used to determine if a
    # trace is generated by a scorer or prediction during evaluation and filter out the scorer
    # traces in the UI. This supposed to be immutable, but we use tag because we can only set this
    # after the scorer is executed, which is not possible with trace metadata.
    SOURCE_SCORER_NAME = "mlflow.trace.sourceScorer"
    # Store a list of linked prompt versions in JSON format
    # Structure: [{"name": "prompt_name", "version": "version"}]
    LINKED_PROMPTS = "mlflow.linkedPrompts"


class TokenUsageKey:
    """Key for the token usage information in the `mlflow.chat.tokenUsage` span attribute."""

    INPUT_TOKENS = "input_tokens"
    OUTPUT_TOKENS = "output_tokens"
    TOTAL_TOKENS = "total_tokens"
    CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens"
    CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens"

    @classmethod
    def all_keys(cls):
        return [
            cls.INPUT_TOKENS,
            cls.OUTPUT_TOKENS,
            cls.TOTAL_TOKENS,
            cls.CACHE_READ_INPUT_TOKENS,
            cls.CACHE_CREATION_INPUT_TOKENS,
        ]

    @classmethod
    def cache_keys(cls):
        return [cls.CACHE_READ_INPUT_TOKENS, cls.CACHE_CREATION_INPUT_TOKENS]


class CostKey:
    """Key for the cost information in the `mlflow.llm.cost` span attribute."""

    INPUT_COST = "input_cost"
    OUTPUT_COST = "output_cost"
    TOTAL_COST = "total_cost"


class TraceSizeStatsKey:
    TOTAL_SIZE_BYTES = "total_size_bytes"
    NUM_SPANS = "num_spans"
    MAX_SPAN_SIZE_BYTES = "max"
    P25_SPAN_SIZE_BYTES = "p25"
    P50_SPAN_SIZE_BYTES = "p50"
    P75_SPAN_SIZE_BYTES = "p75"


# A set of reserved attribute keys
class SpanAttributeKey:
    EXPERIMENT_ID = "mlflow.experimentId"
    REQUEST_ID = "mlflow.traceRequestId"
    INPUTS = "mlflow.spanInputs"
    OUTPUTS = "mlflow.spanOutputs"
    SPAN_TYPE = "mlflow.spanType"
    # Severity level of the span (one of the SpanLogLevel members). Absent
    # means the span was not classified.
    LOG_LEVEL = "mlflow.spanLogLevel"
    FUNCTION_NAME = "mlflow.spanFunctionName"
    START_TIME_NS = "mlflow.spanStartTimeNs"
    CHAT_TOOLS = "mlflow.chat.tools"
    # This attribute is used to store token usage information from LLM responses.
    # Stored in {"input_tokens": int, "output_tokens": int, "total_tokens": int,
    #   "cache_read_input_tokens"?: int, "cache_creation_input_tokens"?: int} format.
    CHAT_USAGE = "mlflow.chat.tokenUsage"
    # This attribute stores cost information calculated from token usage and model pricing.
    # Stored in {"input_cost": float, "output_cost": float, "total_cost": float} format (USD).
    LLM_COST = "mlflow.llm.cost"
    # This attribute stores the model name extracted from span inputs/attributes.
    MODEL = "mlflow.llm.model"
    MODEL_PROVIDER = "mlflow.llm.provider"
    # This attribute indicates which flavor/format generated the LLM span. This is
    # used by downstream (e.g., UI) to determine the message format for parsing.
    MESSAGE_FORMAT = "mlflow.message.format"
    # This attribute is used to populate `intermediate_outputs` property of a trace data
    # representing intermediate outputs of the trace. This attribute is not empty only on
    # the root span of a trace created by the `mlflow.log_trace` API. The `intermediate_outputs`
    # property of the normal trace is generated by the outputs of non-root spans.
    INTERMEDIATE_OUTPUTS = "mlflow.trace.intermediate_outputs"
    # This attribute is used to store prompt version information when load_prompt is called
    # within an active span. Stored as a JSON list of {"name": "...", "version": "..."} objects,
    # same format as LINKED_PROMPTS_TAG_KEY in traces.
    LINKED_PROMPTS = "mlflow.linkedPrompts"
    # This attribute stores the trace ID of the linked gateway trace, used when a gateway
    # endpoint is called by a traced agent via distributed tracing (traceparent header).
    LINKED_GATEWAY_TRACE_ID = "mlflow.gateway.linkedTraceId"

    # User and session IDs copied from trace metadata to root span attributes for OTLP export.
    # Following OTel semantic conventions:
    # https://opentelemetry.io/docs/specs/semconv/registry/attributes/user/#user-id
    # https://opentelemetry.io/docs/specs/semconv/registry/attributes/session/#session-id
    USER_ID = "user.id"
    SESSION_ID = "session.id"
    # Prefix for per-tag span attributes emitted by OtelSpanProcessor so user-defined tags
    # survive OTLP export and can be restored to SqlTraceTag rows on the server side.
    # Each attribute is keyed as "mlflow.traceTag.<tag_key>" with the plain string value.
    TRACE_TAG_PREFIX = "mlflow.traceTag."


class TraceExperimentTagKey:
    ARCHIVAL_RETENTION = "mlflow.trace.archivalRetention"
    ARCHIVE_NOW = "mlflow.trace.archiveNow"


class TraceArchivalFailureReason(str, Enum):
    MALFORMED_TRACE = "MALFORMED_TRACE"
    UNSUPPORTED_ARCHIVE_REPOSITORY = "UNSUPPORTED_ARCHIVE_REPOSITORY"


class AssessmentMetadataKey:
    # When the assessment is generated by an eval run, log the run ID here.
    SOURCE_RUN_ID = "mlflow.assessment.sourceRunId"
    # Total LLM cost spent for generating the feedback (llm-as-a-judge).
    JUDGE_COST = "mlflow.assessment.judgeCost"
    # Token counts for the judge LLM call.
    JUDGE_INPUT_TOKENS = "mlflow.assessment.judgeInputTokens"
    JUDGE_OUTPUT_TOKENS = "mlflow.assessment.judgeOutputTokens"
    # When the scorer generates a trace for assessment scoring, log the trace ID here.
    SCORER_TRACE_ID = "mlflow.assessment.scorerTraceId"
    # When the assessment is generated by online scoring, log the session ID here.
    ONLINE_SCORING_SESSION_ID = "mlflow.assessment.onlineScoringSessionId"


# All storage backends are guaranteed to support request_metadata key/value up to 250 characters
MAX_CHARS_IN_TRACE_INFO_METADATA = 250
# All storage backends are guaranteed to support tag keys up to 250 characters,
# values up to 4096 characters
MAX_CHARS_IN_TRACE_INFO_TAGS_KEY = 250
MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE = 4096
TRUNCATION_SUFFIX = "..."

TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_DBX = 10000
TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_OSS = 1000

# Trace request ID must have the prefix "tr-" appended to the OpenTelemetry trace ID
TRACE_REQUEST_ID_PREFIX = "tr-"
# Trace ID V4 format starts with "trace:/" in the format of "trace:/<location>/<trace_id>"
TRACE_ID_V4_PREFIX = "trace:/"

# Schema version of traces and spans.
TRACE_SCHEMA_VERSION = 3

# Key for the trace schema version in the trace. This key is also used in
# Databricks model serving to be careful when modifying it.
TRACE_SCHEMA_VERSION_KEY = "mlflow.trace_schema.version"


STREAM_CHUNK_EVENT_NAME_FORMAT = "mlflow.chunk.item.{index}"
STREAM_CHUNK_EVENT_VALUE_KEY = "mlflow.chunk.value"


# Key for Databricks model serving options to return the trace in the response
DATABRICKS_OPTIONS_KEY = "databricks_options"
RETURN_TRACE_OPTION_KEY = "return_trace"
DATABRICKS_OUTPUT_KEY = "databricks_output"

# Assessment constants
ASSESSMENT_ID_PREFIX = "a-"


# The location of the spans in the trace.
# This is used to determine where the spans are stored when exporting.
class SpansLocation(str, Enum):
    TRACKING_STORE = "TRACKING_STORE"
    ARTIFACT_REPO = "ARTIFACT_REPO"
    ARCHIVE_REPO = "ARCHIVE_REPO"


# Path to the notebook trace renderer directory
TRACE_RENDERER_ASSET_PATH = "/static-files/lib/notebook-trace-renderer"


class TraceMetricKey:
    """
    Keys for metrics on traces view type.
    """

    TRACE_COUNT = "trace_count"
    SESSION_COUNT = "session_count"
    LATENCY = "latency"
    INPUT_TOKENS = "input_tokens"
    OUTPUT_TOKENS = "output_tokens"
    TOTAL_TOKENS = "total_tokens"
    CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens"
    CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens"

    @classmethod
    def token_usage_keys(cls) -> list[str]:
        return [
            cls.INPUT_TOKENS,
            cls.OUTPUT_TOKENS,
            cls.TOTAL_TOKENS,
            cls.CACHE_READ_INPUT_TOKENS,
            cls.CACHE_CREATION_INPUT_TOKENS,
        ]


class TraceMetricDimensionKey:
    """
    Dimensions for metrics on traces view type.
    """

    TRACE_NAME = "trace_name"
    TRACE_STATUS = "trace_status"


class SpanMetricKey:
    """
    Keys for metrics on spans view type.
    """

    SPAN_COUNT = "span_count"
    LATENCY = "latency"
    INPUT_COST = "input_cost"
    OUTPUT_COST = "output_cost"
    TOTAL_COST = "total_cost"

    @classmethod
    def cost_keys(cls) -> list[str]:
        return [cls.INPUT_COST, cls.OUTPUT_COST, cls.TOTAL_COST]


class SpanMetricDimensionKey:
    """
    Dimensions for metrics on spans view type.
    """

    SPAN_NAME = "span_name"
    SPAN_TYPE = "span_type"
    SPAN_STATUS = "span_status"
    SPAN_MODEL_NAME = "span_model_name"
    SPAN_MODEL_PROVIDER = "span_model_provider"


class AssessmentMetricKey:
    """
    Keys for metrics on assessments view type.
    """

    ASSESSMENT_COUNT = "assessment_count"
    ASSESSMENT_VALUE = "assessment_value"


class AssessmentMetricDimensionKey:
    """
    Dimensions for metrics on assessments view type.
    """

    ASSESSMENT_NAME = "assessment_name"
    ASSESSMENT_VALUE = "assessment_value"


class TraceMetricSearchKey:
    """
    Search key for trace metrics view type.
    VIEW_TYPE must be the prefix of the search string
    e.g. "trace.<entity> = <value>" or "trace.<entity>.<key> = <value>"
    """

    VIEW_TYPE = "trace"
    STATUS = "status"
    TAG = "tag"
    METADATA = "metadata"

    @classmethod
    def entity_to_key_requirement(cls) -> dict[str, bool]:
        """
        Mapping of entity to a boolean indicating if it requires a key
        For example, "tag" requires a key: "trace.tag.<key> = <value>"
        "status" does not require a key: "trace.status = <value>"
        """
        return {
            cls.STATUS: False,
            cls.TAG: True,
            cls.METADATA: True,
        }


class SpanMetricSearchKey:
    """
    Search key for span metrics view type.
    VIEW_TYPE must be the prefix of the search string
    e.g. "span.<entity> = <value>"
    """

    VIEW_TYPE = "span"
    NAME = "name"
    STATUS = "status"
    TYPE = "type"

    @classmethod
    def entity_to_key_requirement(cls) -> dict[str, bool]:
        """
        Mapping of entity to a boolean indicating if it requires a key
        For example, "name" does not require a key: "span.name = <value>"
        """
        return {
            cls.NAME: False,
            cls.STATUS: False,
            cls.TYPE: False,
        }


class AssessmentMetricSearchKey:
    """
    Search key for assessment metrics view type.
    VIEW_TYPE must be the prefix of the search string
    e.g. "assessment.<entity> = <value>"
    """

    VIEW_TYPE = "assessment"
    NAME = "name"
    TYPE = "type"

    @classmethod
    def entity_to_key_requirement(cls) -> dict[str, bool]:
        """
        Mapping of entity to a boolean indicating if it requires a key
        For example, "name" does not require a key: "assessment.name = <value>"
        """
        return {
            cls.NAME: False,
            cls.TYPE: False,
        }


# OpenTelemetry GenAI Semantic Convention attribute keys.
# https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/
class GenAiSemconvKey:
    CONVERSATION_ID = "gen_ai.conversation.id"
    OPERATION_NAME = "gen_ai.operation.name"
    REQUEST_MODEL = "gen_ai.request.model"
    RESPONSE_MODEL = "gen_ai.response.model"
    RESPONSE_ID = "gen_ai.response.id"
    PROVIDER_NAME = "gen_ai.provider.name"
    USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
    USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
    INPUT_MESSAGES = "gen_ai.input.messages"
    OUTPUT_MESSAGES = "gen_ai.output.messages"
    SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions"
    REQUEST_TEMPERATURE = "gen_ai.request.temperature"
    REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens"
    REQUEST_TOP_P = "gen_ai.request.top_p"
    REQUEST_STOP_SEQUENCES = "gen_ai.request.stop_sequences"
    RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons"
    TOOL_DEFINITIONS = "gen_ai.tool.definitions"
    TOOL_CALL_ARGUMENTS = "gen_ai.tool.call.arguments"
    TOOL_CALL_RESULT = "gen_ai.tool.call.result"