# 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"
    # 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"


class TraceTagKey:
    TRACE_NAME = "mlflow.traceName"
    EVAL_REQUEST_ID = "eval.requestId"


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"

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


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"
    FUNCTION_NAME = "mlflow.spanFunctionName"
    START_TIME_NS = "mlflow.spanStartTimeNs"
    # these attributes are for standardized chat messages and tool definitions
    # in CHAT_MODEL and LLM spans. they are used for rendering the rich chat
    # display in the trace UI, as well as downstream consumers of trace data
    # such as evaluation
    CHAT_MESSAGES = "mlflow.chat.messages"
    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} format.
    CHAT_USAGE = "mlflow.chat.tokenUsage"
    # 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"


class AssessmentMetadataKey:
    SOURCE_RUN_ID = "mlflow.assessment.sourceRunId"


# 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-"

# 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"
