mlflow--mlflow/tests/tracing/utils/test_utils.py

import json
import logging
from unittest import mock
from unittest.mock import Mock, patch

import litellm
import pytest
from opentelemetry import trace as trace_api
from pydantic import ValidationError

import mlflow
from mlflow.entities import (
    LiveSpan,
    SpanType,
)
from mlflow.entities.span import SpanType
from mlflow.entities.trace_location import UCSchemaLocation
from mlflow.exceptions import MlflowException
from mlflow.tracing import set_span_chat_tools
from mlflow.tracing.constant import (
    TRACE_ID_V4_PREFIX,
    CostKey,
    SpanAttributeKey,
    TokenUsageKey,
)
from mlflow.tracing.utils import (
    _calculate_percentile,
    aggregate_cost_from_spans,
    aggregate_usage_from_spans,
    calculate_cost_by_model_and_token_usage,
    capture_function_input_args,
    construct_full_inputs,
    dump_span_attribute_value,
    encode_span_id,
    encode_trace_id,
    generate_trace_id_v4,
    generate_trace_id_v4_from_otel_trace_id,
    get_active_spans_table_name,
    get_otel_attribute,
    maybe_get_request_id,
    parse_trace_id_v4,
)
from mlflow.version import IS_TRACING_SDK_ONLY

from tests.tracing.helper import create_mock_otel_span


def test_capture_function_input_args_does_not_raise():
    # Exception during inspecting inputs: trace should be logged without inputs field
    with patch("inspect.signature", side_effect=ValueError("Some error")) as mock_input_args:
        args = capture_function_input_args(lambda: None, (), {})

    assert args is None
    assert mock_input_args.call_count > 0


def test_duplicate_span_names():
    span_names = ["red", "red", "blue", "red", "green", "blue"]

    spans = [
        LiveSpan(create_mock_otel_span("trace_id", span_id=i, name=span_name), trace_id="tr-123")
        for i, span_name in enumerate(span_names)
    ]

    assert [span.name for span in spans] == span_names
    # Check if the span order is preserved
    assert [span.span_id for span in spans] == [encode_span_id(i) for i in [0, 1, 2, 3, 4, 5]]


def test_aggregate_usage_from_spans():
    spans = [
        LiveSpan(create_mock_otel_span("trace_id", span_id=i, name=f"span_{i}"), trace_id="tr-123")
        for i in range(3)
    ]
    spans[0].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 10,
            TokenUsageKey.OUTPUT_TOKENS: 20,
            TokenUsageKey.TOTAL_TOKENS: 30,
        },
    )
    spans[1].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {TokenUsageKey.OUTPUT_TOKENS: 15, TokenUsageKey.TOTAL_TOKENS: 15},
    )
    spans[2].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 5,
            TokenUsageKey.OUTPUT_TOKENS: 10,
            TokenUsageKey.TOTAL_TOKENS: 15,
        },
    )

    usage = aggregate_usage_from_spans(spans)
    assert usage == {
        TokenUsageKey.INPUT_TOKENS: 15,
        TokenUsageKey.OUTPUT_TOKENS: 45,
        TokenUsageKey.TOTAL_TOKENS: 60,
    }


def test_aggregate_usage_from_spans_skips_descendant_usage():
    spans = [
        LiveSpan(create_mock_otel_span("trace_id", span_id=1, name="root"), trace_id="tr-123"),
        LiveSpan(
            create_mock_otel_span("trace_id", span_id=2, name="child", parent_id=1),
            trace_id="tr-123",
        ),
        LiveSpan(
            create_mock_otel_span("trace_id", span_id=3, name="grandchild", parent_id=2),
            trace_id="tr-123",
        ),
        LiveSpan(
            create_mock_otel_span("trace_id", span_id=4, name="independent"), trace_id="tr-123"
        ),
    ]

    spans[0].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 10,
            TokenUsageKey.OUTPUT_TOKENS: 20,
            TokenUsageKey.TOTAL_TOKENS: 30,
        },
    )

    spans[2].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 5,
            TokenUsageKey.OUTPUT_TOKENS: 10,
            TokenUsageKey.TOTAL_TOKENS: 15,
        },
    )

    spans[3].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 3,
            TokenUsageKey.OUTPUT_TOKENS: 6,
            TokenUsageKey.TOTAL_TOKENS: 9,
        },
    )

    usage = aggregate_usage_from_spans(spans)

    assert usage == {
        TokenUsageKey.INPUT_TOKENS: 13,
        TokenUsageKey.OUTPUT_TOKENS: 26,
        TokenUsageKey.TOTAL_TOKENS: 39,
    }


def _deep_chain(spans, start_id, length, parent_id, name):
    # Append a chain of `length` spans (no usage) under `parent_id`, returning the id of
    # the deepest span. Used to push traversal past the recursion limit.
    for i in range(length):
        span_id = start_id + i
        spans.append(
            LiveSpan(
                create_mock_otel_span(
                    "trace_id", span_id=span_id, name=f"{name}_{i}", parent_id=parent_id
                ),
                trace_id="tr-123",
            )
        )
        parent_id = span_id
    return parent_id


def test_aggregate_usage_from_spans_deep_tree_aggregates_leaves():
    # A deep backbone (no usage) that exceeds the recursion limit, ending in a fan of
    # sibling leaves that each carry usage. None of the leaves is an ancestor of another,
    # so aggregation must SUM all of them — the fix must survive the depth AND still
    # aggregate. Regression test for #24344.
    spans = [LiveSpan(create_mock_otel_span("trace_id", span_id=1, name="root"), trace_id="tr-123")]
    deepest = _deep_chain(spans, start_id=2, length=1100, parent_id=1, name="backbone")

    num_leaves = 5
    for j in range(num_leaves):
        leaf = LiveSpan(
            create_mock_otel_span(
                "trace_id", span_id=10_000 + j, name=f"leaf_{j}", parent_id=deepest
            ),
            trace_id="tr-123",
        )
        leaf.set_attribute(
            SpanAttributeKey.CHAT_USAGE,
            {
                TokenUsageKey.INPUT_TOKENS: 2,
                TokenUsageKey.OUTPUT_TOKENS: 3,
                TokenUsageKey.TOTAL_TOKENS: 5,
            },
        )
        spans.append(leaf)

    # All 5 sibling leaves are summed: 5 * {2, 3, 5}.
    assert aggregate_usage_from_spans(spans) == {
        TokenUsageKey.INPUT_TOKENS: 10,
        TokenUsageKey.OUTPUT_TOKENS: 15,
        TokenUsageKey.TOTAL_TOKENS: 25,
    }


def test_aggregate_usage_from_spans_deep_tree_sums_and_skips_descendants():
    # A deep tree where usage lives on two independent branches (both counted and summed)
    # and also on a descendant of a data-bearing span (skipped). Verifies that both real
    # summation AND the anti-double-counting invariant survive past the recursion limit.
    spans = [LiveSpan(create_mock_otel_span("trace_id", span_id=1, name="root"), trace_id="tr-123")]

    # Branch 1: a deep chain off the root. Its top node carries usage (counted); its
    # deepest node also carries usage (a descendant of the top -> must be skipped).
    _deep_chain(spans, start_id=100, length=1100, parent_id=1, name="b1")
    branch1_top = spans[1]  # first span appended by the chain, child of root
    branch1_top.set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 100,
            TokenUsageKey.OUTPUT_TOKENS: 200,
            TokenUsageKey.TOTAL_TOKENS: 300,
        },
    )
    spans[-1].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 999,
            TokenUsageKey.OUTPUT_TOKENS: 999,
            TokenUsageKey.TOTAL_TOKENS: 999,
        },
    )

    # Branch 2: an independent node off the root (not a descendant of branch 1) -> counted.
    branch2 = LiveSpan(
        create_mock_otel_span("trace_id", span_id=5000, name="b2", parent_id=1),
        trace_id="tr-123",
    )
    branch2.set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 1,
            TokenUsageKey.OUTPUT_TOKENS: 2,
            TokenUsageKey.TOTAL_TOKENS: 3,
        },
    )
    spans.append(branch2)

    # branch1_top (100/200/300) + branch2 (1/2/3); the deep descendant of branch1 is skipped.
    assert aggregate_usage_from_spans(spans) == {
        TokenUsageKey.INPUT_TOKENS: 101,
        TokenUsageKey.OUTPUT_TOKENS: 202,
        TokenUsageKey.TOTAL_TOKENS: 303,
    }


def test_aggregate_usage_from_spans_with_cached_tokens():
    spans = [
        LiveSpan(create_mock_otel_span("trace_id", span_id=i, name=f"span_{i}"), trace_id="tr-123")
        for i in range(3)
    ]
    spans[0].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 100,
            TokenUsageKey.OUTPUT_TOKENS: 50,
            TokenUsageKey.TOTAL_TOKENS: 150,
            TokenUsageKey.CACHE_READ_INPUT_TOKENS: 80,
        },
    )
    spans[1].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 200,
            TokenUsageKey.OUTPUT_TOKENS: 100,
            TokenUsageKey.TOTAL_TOKENS: 300,
            TokenUsageKey.CACHE_READ_INPUT_TOKENS: 120,
            TokenUsageKey.CACHE_CREATION_INPUT_TOKENS: 50,
        },
    )
    # span without cached tokens
    spans[2].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 10,
            TokenUsageKey.OUTPUT_TOKENS: 5,
            TokenUsageKey.TOTAL_TOKENS: 15,
        },
    )

    usage = aggregate_usage_from_spans(spans)
    assert usage == {
        TokenUsageKey.INPUT_TOKENS: 310,
        TokenUsageKey.OUTPUT_TOKENS: 155,
        TokenUsageKey.TOTAL_TOKENS: 465,
        TokenUsageKey.CACHE_READ_INPUT_TOKENS: 200,
        TokenUsageKey.CACHE_CREATION_INPUT_TOKENS: 50,
    }


def test_aggregate_usage_from_spans_without_cached_tokens_omits_keys():
    spans = [
        LiveSpan(create_mock_otel_span("trace_id", span_id=0, name="span_0"), trace_id="tr-123")
    ]
    spans[0].set_attribute(
        SpanAttributeKey.CHAT_USAGE,
        {
            TokenUsageKey.INPUT_TOKENS: 10,
            TokenUsageKey.OUTPUT_TOKENS: 5,
            TokenUsageKey.TOTAL_TOKENS: 15,
        },
    )

    usage = aggregate_usage_from_spans(spans)
    assert usage == {
        TokenUsageKey.INPUT_TOKENS: 10,
        TokenUsageKey.OUTPUT_TOKENS: 5,
        TokenUsageKey.TOTAL_TOKENS: 15,
    }
    # Cached keys should not be present
    assert TokenUsageKey.CACHE_READ_INPUT_TOKENS not in usage
    assert TokenUsageKey.CACHE_CREATION_INPUT_TOKENS not in usage


def test_aggregate_cost_from_spans():
    spans = [
        LiveSpan(create_mock_otel_span("trace_id", span_id=i, name=f"span_{i}"), trace_id="tr-123")
        for i in range(3)
    ]
    spans[0].set_attribute(
        SpanAttributeKey.LLM_COST,
        {
            CostKey.INPUT_COST: 10,
            CostKey.OUTPUT_COST: 20,
            CostKey.TOTAL_COST: 30,
        },
    )
    spans[1].set_attribute(
        SpanAttributeKey.LLM_COST,
        {CostKey.OUTPUT_COST: 15, CostKey.TOTAL_COST: 15},
    )
    spans[2].set_attribute(
        SpanAttributeKey.LLM_COST,
        {
            CostKey.INPUT_COST: 5,
            CostKey.OUTPUT_COST: 10,
            CostKey.TOTAL_COST: 15,
        },
    )

    cost = aggregate_cost_from_spans(spans)
    assert cost == {
        CostKey.INPUT_COST: 15,
        CostKey.OUTPUT_COST: 45,
        CostKey.TOTAL_COST: 60,
    }


def test_aggregate_cost_from_spans_skips_descendant_cost():
    spans = [
        LiveSpan(create_mock_otel_span("trace_id", span_id=1, name="root"), trace_id="tr-123"),
        LiveSpan(
            create_mock_otel_span("trace_id", span_id=2, name="child", parent_id=1),
            trace_id="tr-123",
        ),
        LiveSpan(
            create_mock_otel_span("trace_id", span_id=3, name="grandchild", parent_id=2),
            trace_id="tr-123",
        ),
        LiveSpan(
            create_mock_otel_span("trace_id", span_id=4, name="independent"), trace_id="tr-123"
        ),
    ]

    spans[0].set_attribute(
        SpanAttributeKey.LLM_COST,
        {
            CostKey.INPUT_COST: 10,
            CostKey.OUTPUT_COST: 20,
            CostKey.TOTAL_COST: 30,
        },
    )

    spans[2].set_attribute(
        SpanAttributeKey.LLM_COST,
        {
            CostKey.INPUT_COST: 5,
            CostKey.OUTPUT_COST: 10,
            CostKey.TOTAL_COST: 15,
        },
    )

    spans[3].set_attribute(
        SpanAttributeKey.LLM_COST,
        {
            CostKey.INPUT_COST: 3,
            CostKey.OUTPUT_COST: 6,
            CostKey.TOTAL_COST: 9,
        },
    )

    cost = aggregate_cost_from_spans(spans)

    assert cost == {
        CostKey.INPUT_COST: 13,
        CostKey.OUTPUT_COST: 26,
        CostKey.TOTAL_COST: 39,
    }


def test_maybe_get_request_id():
    assert maybe_get_request_id(is_evaluate=True) is None

    try:
        from mlflow.pyfunc.context import Context, set_prediction_context
    except ImportError:
        pytest.skip("Skipping the rest of tests as mlflow.pyfunc module is not available.")

    with set_prediction_context(Context(request_id="eval", is_evaluate=True)):
        assert maybe_get_request_id(is_evaluate=True) == "eval"

    with set_prediction_context(Context(request_id="non_eval", is_evaluate=False)):
        assert maybe_get_request_id(is_evaluate=True) is None


def test_set_chat_tools_validation():
    tools = [
        {
            "type": "unsupported_function",
            "unsupported_function": {
                "name": "test",
            },
        }
    ]

    @mlflow.trace(span_type=SpanType.CHAT_MODEL)
    def dummy_call(tools):
        span = mlflow.get_current_active_span()
        set_span_chat_tools(span, tools)
        return None

    with pytest.raises(ValidationError, match="validation error for ChatTool"):
        dummy_call(tools)


@pytest.mark.parametrize(
    ("enum_values", "param_type"),
    [
        ([1, 2, 3, 4, 5], "integer"),
        (["option1", "option2", "option3"], "string"),
        ([1.1, 2.5, 3.7], "number"),
        ([True, False], "boolean"),
        (["mixed", 42, True, 3.14], "string"),  # Mixed types with string base type
    ],
)
def test_openai_parse_tools_enum_validation(enum_values, param_type):
    from mlflow.openai.utils.chat_schema import _parse_tools

    # Simulate the exact OpenAI autologging input that was failing
    openai_inputs = {
        "tools": [
            {
                "type": "function",
                "function": {
                    "name": "select_option",
                    "description": "Select an option from the given choices",
                    "parameters": {
                        "type": "object",
                        "properties": {"option": {"type": param_type, "enum": enum_values}},
                        "required": ["option"],
                    },
                },
            }
        ]
    }

    # This should not raise a ValidationError - tests the actual failing code path
    parsed_tools = _parse_tools(openai_inputs)
    assert len(parsed_tools) == 1
    assert parsed_tools[0].function.name == "select_option"
    assert parsed_tools[0].function.parameters.properties["option"].enum == enum_values


def test_construct_full_inputs_simple_function():
    def func(a, b, c=3, d=4, **kwargs):
        pass

    result = construct_full_inputs(func, 1, 2)
    assert result == {"a": 1, "b": 2}

    result = construct_full_inputs(func, 1, 2, c=30)
    assert result == {"a": 1, "b": 2, "c": 30}

    result = construct_full_inputs(func, 1, 2, c=30, d=40, e=50)
    assert result == {"a": 1, "b": 2, "c": 30, "d": 40, "kwargs": {"e": 50}}

    def no_args_func():
        pass

    result = construct_full_inputs(no_args_func)
    assert result == {}

    class TestClass:
        def func(self, a, b, c=3, d=4, **kwargs):
            pass

    result = construct_full_inputs(TestClass().func, 1, 2)
    assert result == {"a": 1, "b": 2}


def test_calculate_percentile():
    # Test empty list
    assert _calculate_percentile([], 0.5) == 0.0

    # Test single element
    assert _calculate_percentile([100], 0.25) == 100
    assert _calculate_percentile([100], 0.5) == 100
    assert _calculate_percentile([100], 0.75) == 100

    # Test two elements
    assert _calculate_percentile([10, 20], 0.0) == 10
    assert _calculate_percentile([10, 20], 0.5) == 15  # Linear interpolation
    assert _calculate_percentile([10, 20], 1.0) == 20

    # Test odd number of elements
    data = [10, 20, 30, 40, 50]
    assert _calculate_percentile(data, 0.0) == 10
    assert _calculate_percentile(data, 0.25) == 20
    assert _calculate_percentile(data, 0.5) == 30  # Median
    assert _calculate_percentile(data, 0.75) == 40
    assert _calculate_percentile(data, 1.0) == 50

    # Test even number of elements
    data = [10, 20, 30, 40]
    assert _calculate_percentile(data, 0.0) == 10
    assert _calculate_percentile(data, 0.25) == 17.5  # Between 10 and 20
    assert _calculate_percentile(data, 0.5) == 25  # Between 20 and 30
    assert _calculate_percentile(data, 0.75) == 32.5  # Between 30 and 40
    assert _calculate_percentile(data, 1.0) == 40

    # Test with larger dataset
    data = list(range(1, 101))  # 1 to 100
    assert _calculate_percentile(data, 0.25) == 25.75
    assert _calculate_percentile(data, 0.5) == 50.5


def test_parse_trace_id_v4():
    test_trace_id = "tr-original-trace-123"

    v4_id_uc_schema = f"{TRACE_ID_V4_PREFIX}catalog.schema/{test_trace_id}"
    location, parsed_id = parse_trace_id_v4(v4_id_uc_schema)
    assert location == "catalog.schema"
    assert parsed_id == test_trace_id

    v4_id_experiment = f"{TRACE_ID_V4_PREFIX}experiment_id/{test_trace_id}"
    location, parsed_id = parse_trace_id_v4(v4_id_experiment)
    assert location == "experiment_id"
    assert parsed_id == test_trace_id

    location, parsed_id = parse_trace_id_v4(test_trace_id)
    assert location is None
    assert parsed_id == test_trace_id


def test_parse_trace_id_v4_invalid_format():
    with pytest.raises(MlflowException, match="Invalid trace ID format"):
        parse_trace_id_v4(f"{TRACE_ID_V4_PREFIX}123")

    with pytest.raises(MlflowException, match="Invalid trace ID format"):
        parse_trace_id_v4(f"{TRACE_ID_V4_PREFIX}123/")

    with pytest.raises(MlflowException, match="Invalid trace ID format"):
        parse_trace_id_v4(f"{TRACE_ID_V4_PREFIX}catalog.schema/../invalid-trace-id")

    with pytest.raises(MlflowException, match="Invalid trace ID format"):
        parse_trace_id_v4(f"{TRACE_ID_V4_PREFIX}catalog.schema/invalid-trace-id/invalid-format")


def test_get_otel_attribute_existing_attribute():
    # Create a mock span with attributes
    span = Mock(spec=trace_api.Span)
    span.attributes = {
        "test_key": json.dumps({"data": "value"}),
        "string_key": json.dumps("simple_string"),
        "number_key": json.dumps(42),
        "boolean_key": json.dumps(True),
        "list_key": json.dumps([1, 2, 3]),
    }

    # Test various data types
    result = get_otel_attribute(span, "test_key")
    assert result == {"data": "value"}

    result = get_otel_attribute(span, "string_key")
    assert result == "simple_string"

    result = get_otel_attribute(span, "number_key")
    assert result == 42

    result = get_otel_attribute(span, "boolean_key")
    assert result is True

    result = get_otel_attribute(span, "list_key")
    assert result == [1, 2, 3]


def test_get_otel_attribute_missing_attribute():
    # Create a mock span with empty attributes
    span = Mock(spec=trace_api.Span)
    span.attributes = {}

    result = get_otel_attribute(span, "nonexistent_key")
    assert result is None


def test_get_otel_attribute_none_attribute():
    # Create a mock span where attributes.get() returns None
    span = Mock(spec=trace_api.Span)
    span.attributes = Mock()
    span.attributes.get.return_value = None

    result = get_otel_attribute(span, "any_key")
    assert result is None


def test_get_otel_attribute_invalid_json():
    # Create a mock span with invalid JSON
    span = Mock(spec=trace_api.Span)
    span.attributes = {
        "invalid_json": "not valid json {",
        "empty_string": "",
    }

    result = get_otel_attribute(span, "invalid_json")
    assert result is None

    result = get_otel_attribute(span, "empty_string")
    assert result is None


def test_get_otel_attribute_non_string_attribute():
    # In some edge cases, attributes might contain non-string values
    span = Mock(spec=trace_api.Span)
    span.attributes = {
        "number_value": 123,  # Not a JSON string
        "boolean_value": True,  # Not a JSON string
    }

    # These should fail gracefully and return None
    result = get_otel_attribute(span, "number_value")
    assert result is None

    result = get_otel_attribute(span, "boolean_value")
    assert result is None


def test_generate_trace_id_v4_with_uc_schema():
    span = create_mock_otel_span(trace_id=12345, span_id=1)
    uc_schema = "catalog.schema"

    with mock.patch(
        "mlflow.tracing.utils.construct_trace_id_v4", return_value="trace:/catalog.schema/abc123"
    ) as mock_construct:
        result = generate_trace_id_v4(span, uc_schema)

        mock_construct.assert_called_once_with(uc_schema, mock.ANY)
        assert result == "trace:/catalog.schema/abc123"


def test_get_spans_table_name_for_trace_with_destination():
    mock_destination = UCSchemaLocation(catalog_name="catalog", schema_name="schema")

    with mock.patch("mlflow.tracing.provider._MLFLOW_TRACE_USER_DESTINATION") as mock_ctx:
        mock_ctx.get.return_value = mock_destination

        result = get_active_spans_table_name()
        assert result == "catalog.schema.mlflow_experiment_trace_otel_spans"


def test_get_spans_table_name_for_trace_no_destination():
    with mock.patch("mlflow.tracing.provider._MLFLOW_TRACE_USER_DESTINATION") as mock_ctx:
        mock_ctx.get.return_value = None

        result = get_active_spans_table_name()
        assert result is None


@pytest.mark.skipif(IS_TRACING_SDK_ONLY, reason="mock_litellm_cost cannot affect server-side cost")
@pytest.mark.parametrize("is_databricks", [True, False])
def test_cost_not_computed_client_side(is_databricks, mock_litellm_cost):
    # Mock should_compute_cost_client_side in the span module (where it's bound at import time)
    # rather than is_databricks_uri. Mocking is_databricks_uri captures the reference in
    # mlflow_v3.py during lazy import and causes _export_spans_incrementally to skip spans.
    with (
        mock.patch(
            "mlflow.entities.span.should_compute_cost_client_side", return_value=is_databricks
        ),
        mock.patch(
            "mlflow.tracing.processor.base_mlflow.should_compute_cost_client_side",
            return_value=is_databricks,
        ),
        mock.patch(
            "mlflow.entities.span.set_span_cost_attribute", wraps=lambda span: None
        ) as mock_set_cost,
    ):
        with mlflow.start_span(name="llm_span") as span:
            span.set_attribute(SpanAttributeKey.MODEL, "gpt-5")
            span.set_attribute(
                SpanAttributeKey.CHAT_USAGE,
                {
                    TokenUsageKey.INPUT_TOKENS: 100,
                    TokenUsageKey.OUTPUT_TOKENS: 50,
                    TokenUsageKey.TOTAL_TOKENS: 150,
                },
            )
        # Cost should be computed at server side if not in Databricks
        if is_databricks:
            mock_set_cost.assert_called_once()
        else:
            mock_set_cost.assert_not_called()

    trace = mlflow.get_trace(trace_id=span.trace_id, flush=True)
    # cost should be set
    assert trace.info.cost is not None
    assert CostKey.INPUT_COST in trace.info.cost
    assert CostKey.OUTPUT_COST in trace.info.cost
    assert CostKey.TOTAL_COST in trace.info.cost


def test_generate_trace_id_v4_from_otel_trace_id():
    otel_trace_id = 0x12345678901234567890123456789012
    location = "catalog.schema"

    result = generate_trace_id_v4_from_otel_trace_id(otel_trace_id, location)

    # Verify the format is trace:/<location>/<hex_trace_id>
    assert result.startswith(f"{TRACE_ID_V4_PREFIX}{location}/")

    # Extract and verify the hex trace ID part
    expected_hex_id = encode_trace_id(otel_trace_id)
    assert result == f"{TRACE_ID_V4_PREFIX}{location}/{expected_hex_id}"

    # Verify it can be parsed back
    parsed_location, parsed_id = parse_trace_id_v4(result)
    assert parsed_location == location
    assert parsed_id == expected_hex_id


def test_builtin_cost_fallback_when_litellm_unavailable():
    with mock.patch.dict("sys.modules", {"litellm": None}):
        result = calculate_cost_by_model_and_token_usage(
            "gpt-4o", {"input_tokens": 1000, "output_tokens": 500}
        )
    assert result is not None
    assert result["input_cost"] == pytest.approx(0.0025)
    assert result["output_cost"] == pytest.approx(0.005)
    assert result["total_cost"] == pytest.approx(0.0075)


def test_builtin_cost_fallback_returns_none_for_unknown_model():
    with mock.patch.dict("sys.modules", {"litellm": None}):
        result = calculate_cost_by_model_and_token_usage(
            "unknown-model", {"input_tokens": 100, "output_tokens": 50}
        )
    assert result is None


def test_builtin_cost_fallback_with_cache_tokens():
    with mock.patch.dict("sys.modules", {"litellm": None}):
        result = calculate_cost_by_model_and_token_usage(
            "gpt-4o",
            {
                "input_tokens": 1000,
                "output_tokens": 500,
                "cache_read_input_tokens": 200,
            },
        )
    assert result is not None
    assert result["input_cost"] == pytest.approx(0.00225)


def test_builtin_cost_fallback_with_provider():
    with mock.patch.dict("sys.modules", {"litellm": None}):
        result = calculate_cost_by_model_and_token_usage(
            "gpt-4o",
            {"input_tokens": 1000, "output_tokens": 500},
            model_provider="openai",
        )
    assert result is not None
    assert result["total_cost"] == pytest.approx(0.0075)


@pytest.mark.parametrize("model_provider", ["OpenAI", "OPENAI", "openai"])
def test_builtin_cost_fallback_with_provider_case_insensitive(model_provider):
    with mock.patch.dict("sys.modules", {"litellm": None}):
        result = calculate_cost_by_model_and_token_usage(
            "gpt-4o",
            {"input_tokens": 1000, "output_tokens": 500},
            model_provider=model_provider,
        )
    assert result is not None
    assert result["total_cost"] == pytest.approx(0.0075)


@pytest.mark.parametrize("model_name", ["gateway:/my-endpoint", "endpoints:/my-endpoint"])
def test_cost_skipped_for_internal_routing_uris(model_name):
    result = calculate_cost_by_model_and_token_usage(
        model_name, {"input_tokens": 1000, "output_tokens": 500}
    )
    assert result is None


def test_litellm_provider_list_not_printed_during_cost_calculation(capsys):
    litellm.suppress_debug_info = False

    calculate_cost_by_model_and_token_usage(
        model_name="databricks-claude-sonnet-4-5",
        usage={TokenUsageKey.INPUT_TOKENS: 10, TokenUsageKey.OUTPUT_TOKENS: 5},
    )

    captured = capsys.readouterr()
    assert "Provider List" not in captured.out
    assert litellm.suppress_debug_info is False


def test_litellm_provider_list_printed_when_debug_logging(capsys):
    litellm.suppress_debug_info = True

    _logger = logging.getLogger("mlflow.tracing.utils")
    original_level = _logger.level
    _logger.setLevel(logging.DEBUG)
    try:
        calculate_cost_by_model_and_token_usage(
            model_name="databricks-claude-sonnet-4-5",
            usage={TokenUsageKey.INPUT_TOKENS: 10, TokenUsageKey.OUTPUT_TOKENS: 5},
        )
    finally:
        _logger.setLevel(original_level)

    captured = capsys.readouterr()
    assert "Provider List" in captured.out
    # During the call to calculate cost, suppress was set to False
    # We are asserting that suppress is reset to the original value after
    assert litellm.suppress_debug_info is True


def test_dump_span_attribute_value_handles_circular_reference():
    cyclic = {"name": "run_context"}
    cyclic["self"] = cyclic

    with pytest.raises(ValueError, match="Circular reference detected"):
        json.dumps(cyclic)

    # Must not raise; fall back result is a valid JSON string containing repr(value).
    result = dump_span_attribute_value(cyclic)
    loaded = json.loads(result)
    assert isinstance(loaded, str)
    assert "run_context" in loaded


def test_dump_span_attribute_value_handles_type_error():
    value = {frozenset({"listener"}): "handler"}

    with pytest.raises(TypeError, match="frozenset"):
        json.dumps(value)

    result = dump_span_attribute_value(value)

    # Must not raise; fall back result is a valid JSON string containing repr(value).
    assert result == json.dumps(repr(value), ensure_ascii=False)