import json
from typing import Any

import pytest
from openinference.semconv.trace import (
    DocumentAttributes,
    MessageAttributes,
    MessageContentAttributes,
    OpenInferenceMimeTypeValues,
    OpenInferenceSpanKindValues,
    SpanAttributes,
    ToolAttributes,
    ToolCallAttributes,
)

from phoenix.db.models import Span
from phoenix.server.api.helpers.dataset_helpers import (
    _get_message,
    _merge_assistant_output_items,
    get_dataset_example_input,
    get_dataset_example_output,
    get_experiment_example_output,
)
from phoenix.trace.attributes import unflatten

# DocumentAttributes
DOCUMENT_CONTENT = DocumentAttributes.DOCUMENT_CONTENT
DOCUMENT_ID = DocumentAttributes.DOCUMENT_ID
DOCUMENT_SCORE = DocumentAttributes.DOCUMENT_SCORE

# MessageAttributes
MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
MESSAGE_CONTENTS = MessageAttributes.MESSAGE_CONTENTS
MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON = MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON
MESSAGE_FUNCTION_CALL_NAME = MessageAttributes.MESSAGE_FUNCTION_CALL_NAME
MESSAGE_NAME = MessageAttributes.MESSAGE_NAME
MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS

# MessageContentAttributes (OpenAI Responses API output format)
MESSAGE_CONTENT_TEXT = MessageContentAttributes.MESSAGE_CONTENT_TEXT

# OpenInferenceMimeTypeValues
JSON = OpenInferenceMimeTypeValues.JSON.value
TEXT = OpenInferenceMimeTypeValues.TEXT.value

# OpenInferenceSpanKindValues
CHAIN = OpenInferenceSpanKindValues.CHAIN.value
LLM = OpenInferenceSpanKindValues.LLM.value
RETRIEVER = OpenInferenceSpanKindValues.RETRIEVER.value

# SpanAttributes
INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
INPUT_VALUE = SpanAttributes.INPUT_VALUE
LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
LLM_PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS

# ToolCallAttributes
TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
TOOL_CALL_FUNCTION_NAME = ToolCallAttributes.TOOL_CALL_FUNCTION_NAME


@pytest.mark.parametrize(
    "span, expected_input_value",
    [
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (OPENINFERENCE_SPAN_KIND, LLM),
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (
                            LLM_PROMPT_TEMPLATE_VARIABLES,
                            json.dumps({"variable_name": "variable-value"}),
                        ),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "123"),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_NAME}", "xyz"),
                        (f"{LLM_INPUT_MESSAGES}.1.{MESSAGE_CONTENT}", "user-message"),
                        (f"{LLM_INPUT_MESSAGES}.1.{MESSAGE_ROLE}", "user"),
                        (f"{LLM_INPUT_MESSAGES}.2.{MESSAGE_ROLE}", "assistant"),
                        (f"{LLM_INPUT_MESSAGES}.2.{MESSAGE_FUNCTION_CALL_NAME}", "add"),
                        (
                            f"{LLM_INPUT_MESSAGES}.2.{MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON}",
                            json.dumps({"a": 363, "b": 42}),
                        ),
                        (f"{LLM_INPUT_MESSAGES}.3.{MESSAGE_CONTENT}", "user-message"),
                        (f"{LLM_INPUT_MESSAGES}.3.{MESSAGE_ROLE}", "user"),
                        (f"{LLM_INPUT_MESSAGES}.4.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{LLM_INPUT_MESSAGES}.4.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_NAME}",
                            "multiply",
                        ),
                        (
                            f"{LLM_INPUT_MESSAGES}.4.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            json.dumps({"a": 121, "b": 3}),
                        ),
                        (
                            f"{LLM_INPUT_MESSAGES}.4.{MESSAGE_TOOL_CALLS}.1.{TOOL_CALL_FUNCTION_NAME}",
                            "add",
                        ),
                        (
                            f"{LLM_INPUT_MESSAGES}.4.{MESSAGE_TOOL_CALLS}.1.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            json.dumps({"a": 363, "b": 42}),
                        ),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {
                "messages": [
                    {"content": "123", "role": "assistant", "name": "xyz"},
                    {"content": "user-message", "role": "user"},
                    {
                        "role": "assistant",
                        "tool_calls": [
                            {"function": {"name": "add", "arguments": {"a": 363, "b": 42}}},
                        ],
                    },
                    {"content": "user-message", "role": "user"},
                    {
                        "role": "assistant",
                        "tool_calls": [
                            {"function": {"name": "multiply", "arguments": {"a": 121, "b": 3}}},
                            {"function": {"name": "add", "arguments": {"a": 363, "b": 42}}},
                        ],
                    },
                ],
                "variable_name": "variable-value",
            },
            id="llm-span-with-input-messages-and-prompt-template-variables",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "user-message"),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_ROLE}", "user"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {
                "messages": [{"content": "user-message", "role": "user"}],
            },
            id="llm-span-with-input-messages-and-no-prompt-template-variables",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {"input": "plain-text-input"},
            id="llm-span-with-no-input-messages-and-plain-text-input",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (
                            LLM_PROMPT_TEMPLATE_VARIABLES,
                            json.dumps({"variable_name": "variable-value"}),
                        ),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {
                "input": "plain-text-input",
                "variable_name": "variable-value",
            },
            id="llm-span-with-no-input-messages-and-plain-text-input-with-prompt-template-variables",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, json.dumps({"llm-span-input": "llm-input"})),
                        (INPUT_MIME_TYPE, JSON),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {"llm-span-input": "llm-input"},
            id="llm-span-with-no-input-messages-and-json-input",
        ),
        pytest.param(
            Span(
                span_kind="CHAIN",
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {"input": "plain-text-input"},
            id="chain-span-with-plain-text-input",
        ),
        pytest.param(
            Span(
                span_kind="CHAIN",
                attributes=unflatten(
                    (
                        (INPUT_VALUE, json.dumps({"chain_input": "chain-input"})),
                        (INPUT_MIME_TYPE, JSON),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {"chain_input": "chain-input"},
            id="chain-span-with-json-input",
        ),
    ],
)
def test_get_dataset_example_input(span: Span, expected_input_value: dict[str, Any]) -> None:
    input_value = get_dataset_example_input(span)
    assert expected_input_value == input_value


@pytest.mark.parametrize(
    "message, expected",
    [
        pytest.param(
            unflatten([(MESSAGE_ROLE, "user")]),
            {"role": "user"},
            id="role_only",
        ),
        pytest.param(
            unflatten(
                [
                    (MESSAGE_ROLE, "assistant"),
                    (MESSAGE_CONTENT, "hi"),
                    (MESSAGE_NAME, "bot"),
                ]
            ),
            {"role": "assistant", "content": "hi", "name": "bot"},
            id="content_and_name",
        ),
        pytest.param(
            unflatten(
                [
                    (MESSAGE_ROLE, "assistant"),
                    (MESSAGE_FUNCTION_CALL_NAME, "add"),
                    (MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON, json.dumps({"a": 1, "b": 2})),
                ]
            ),
            {
                "role": "assistant",
                "tool_calls": [
                    {"function": {"name": "add", "arguments": {"a": 1, "b": 2}}},
                ],
            },
            id="function_call_json_string_deserialized",
        ),
        pytest.param(
            unflatten(
                [
                    (MESSAGE_ROLE, "assistant"),
                    (MESSAGE_FUNCTION_CALL_NAME, "add"),
                    (MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON, "not valid json"),
                ]
            ),
            {
                "role": "assistant",
                "tool_calls": [
                    {"function": {"name": "add", "arguments": "not valid json"}},
                ],
            },
            id="function_call_invalid_json_left_as_string",
        ),
        pytest.param(
            unflatten(
                [
                    (MESSAGE_ROLE, "assistant"),
                    (
                        MESSAGE_TOOL_CALLS,
                        [
                            unflatten(
                                [
                                    (TOOL_CALL_FUNCTION_NAME, "multiply"),
                                    (
                                        TOOL_CALL_FUNCTION_ARGUMENTS_JSON,
                                        json.dumps({"x": 10}),
                                    ),
                                ]
                            ),
                        ],
                    ),
                ]
            ),
            {
                "role": "assistant",
                "tool_calls": [
                    {"function": {"name": "multiply", "arguments": {"x": 10}}},
                ],
            },
            id="tool_calls_json_string_deserialized",
        ),
        pytest.param(
            unflatten(
                [
                    (MESSAGE_ROLE, "assistant"),
                    (
                        MESSAGE_TOOL_CALLS,
                        [
                            unflatten(
                                [
                                    (TOOL_CALL_FUNCTION_NAME, "run"),
                                    (
                                        TOOL_CALL_FUNCTION_ARGUMENTS_JSON,
                                        "{broken",
                                    ),
                                ]
                            ),
                        ],
                    ),
                ]
            ),
            {
                "role": "assistant",
                "tool_calls": [
                    {"function": {"name": "run", "arguments": "{broken"}},
                ],
            },
            id="tool_calls_invalid_json_left_as_string",
        ),
        pytest.param(
            unflatten([(MESSAGE_ROLE, "user"), (MESSAGE_TOOL_CALLS, [])]),
            {"role": "user"},
            id="empty_tool_calls_omitted",
        ),
        pytest.param(
            unflatten([(MESSAGE_ROLE, "user"), (MESSAGE_TOOL_CALLS, None)]),
            {"role": "user"},
            id="none_tool_calls_omitted",
        ),
    ],
)
def test_get_message(message: dict[str, Any], expected: dict[str, Any]) -> None:
    assert _get_message(message) == expected


@pytest.mark.parametrize(
    "raw, expected",
    [
        pytest.param(
            [],
            [],
            id="empty_input",
        ),
        pytest.param(
            [{"role": "assistant", "content": "hello"}],
            [{"role": "assistant", "content": "hello"}],
            id="single_assistant_passthrough",
        ),
        pytest.param(
            [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hey"}],
            [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hey"}],
            id="non_assistant_not_merged",
        ),
        pytest.param(
            [
                {"role": "assistant", "content": "I will call the tool."},
                {
                    "role": "assistant",
                    "tool_calls": [
                        {"function": {"name": "get_weather", "arguments": {"loc": "SF"}}}
                    ],
                },
            ],
            [
                {
                    "role": "assistant",
                    "content": "I will call the tool.",
                    "tool_calls": [
                        {"function": {"name": "get_weather", "arguments": {"loc": "SF"}}}
                    ],
                }
            ],
            id="content_and_tool_call_merged",
        ),
        pytest.param(
            [
                {
                    "role": "assistant",
                    "tool_calls": [
                        {"function": {"name": "tool_a", "arguments": {}}, "id": "c1"},
                    ],
                },
                {
                    "role": "assistant",
                    "tool_calls": [
                        {"function": {"name": "tool_b", "arguments": {}}, "id": "c2"},
                    ],
                },
            ],
            [
                {
                    "role": "assistant",
                    "tool_calls": [
                        {"function": {"name": "tool_a", "arguments": {}}, "id": "c1"},
                        {"function": {"name": "tool_b", "arguments": {}}, "id": "c2"},
                    ],
                }
            ],
            id="multiple_tool_calls_merged",
        ),
        pytest.param(
            [
                {"role": "user", "content": "question"},
                {"role": "assistant", "content": "answer"},
                {
                    "role": "assistant",
                    "tool_calls": [{"function": {"name": "search", "arguments": {"q": "x"}}}],
                },
                {"role": "user", "content": "follow-up"},
                {"role": "assistant", "content": "done"},
            ],
            [
                {"role": "user", "content": "question"},
                {
                    "role": "assistant",
                    "content": "answer",
                    "tool_calls": [{"function": {"name": "search", "arguments": {"q": "x"}}}],
                },
                {"role": "user", "content": "follow-up"},
                {"role": "assistant", "content": "done"},
            ],
            id="interleaved_roles_merge_only_within_runs",
        ),
        pytest.param(
            [
                {
                    "role": "assistant",
                    "tool_calls": [{"function": {"name": "fn", "arguments": {}}}],
                },
            ],
            [
                {
                    "role": "assistant",
                    "tool_calls": [{"function": {"name": "fn", "arguments": {}}}],
                },
            ],
            id="tool_calls_only_no_content",
        ),
        pytest.param(
            [{"role": "assistant"}, {"role": "assistant"}],
            [{"role": "assistant"}],
            id="consecutive_empty_assistants_merged",
        ),
        pytest.param(
            [
                {"role": "assistant"},
                {"role": "assistant", "content": "late content"},
            ],
            [{"role": "assistant", "content": "late content"}],
            id="content_from_later_item_adopted",
        ),
    ],
)
def test_merge_assistant_output_items(
    raw: list[dict[str, Any]], expected: list[dict[str, Any]]
) -> None:
    assert _merge_assistant_output_items(raw) == expected  # type: ignore[arg-type]


@pytest.mark.parametrize(
    "span, expected_output_value",
    [
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (
                            LLM_PROMPT_TEMPLATE_VARIABLES,
                            json.dumps({"variable_name": "variable-value"}),
                        ),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "user-message"),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_ROLE}", "user"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "assistant-message"),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {"messages": [{"content": "assistant-message", "role": "assistant"}]},
            id="llm-span-with-output-messages",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "user-message"),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_ROLE}", "user"),
                        # OpenAI Responses API stores output under message.contents.*.message_content.text
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENTS}.0.{MESSAGE_CONTENT_TEXT}",
                            "Responses API output text",
                        ),
                    )
                ),
            ),
            {"messages": [{"content": "Responses API output text", "role": "assistant"}]},
            id="llm-span-with-output-messages-responses-api-format",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENTS}.0.{MESSAGE_CONTENT_TEXT}",
                            "First block",
                        ),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENTS}.1.{MESSAGE_CONTENT_TEXT}",
                            "Second block",
                        ),
                    )
                ),
            ),
            {"messages": [{"content": "First block\n\nSecond block", "role": "assistant"}]},
            id="llm-span-with-output-messages-responses-api-format-multiple-parts",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "user-message"),
                        (f"{LLM_INPUT_MESSAGES}.0.{MESSAGE_ROLE}", "user"),
                    )
                ),
            ),
            {"output": "plain-text-output"},
            id="llm-span-with-no-output-messages-but-with-plain-text-output",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, "plain-text-input"),
                        (INPUT_MIME_TYPE, TEXT),
                        (OUTPUT_VALUE, json.dumps({"llm-span-output": "value"})),
                        (OUTPUT_MIME_TYPE, JSON),
                    )
                ),
            ),
            {"llm-span-output": "value"},
            id="llm-span-with-no-output-messages-and-json-output",
        ),
        pytest.param(
            Span(
                span_kind=RETRIEVER,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, json.dumps({"retriever-input": "retriever-input"})),
                        (INPUT_MIME_TYPE, JSON),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (f"{RETRIEVAL_DOCUMENTS}.0.{DOCUMENT_ID}", "1"),
                        (f"{RETRIEVAL_DOCUMENTS}.0.{DOCUMENT_SCORE}", 0.5),
                        (f"{RETRIEVAL_DOCUMENTS}.0.{DOCUMENT_CONTENT}", "document-content"),
                    )
                ),
            ),
            {"documents": [{"id": "1", "score": 0.5, "content": "document-content"}]},
            id="retriever-span-with-retrieval-documents",
        ),
        pytest.param(
            Span(
                span_kind=RETRIEVER,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, json.dumps({"retriever-input": "retriever-input"})),
                        (INPUT_MIME_TYPE, JSON),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                    )
                ),
            ),
            {"output": "plain-text-output"},
            id="retriever-span-with-plain-text-output-and-no-retrieval-documents",
        ),
        pytest.param(
            Span(
                span_kind=RETRIEVER,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, json.dumps({"retriever-input": "retriever-input"})),
                        (INPUT_MIME_TYPE, JSON),
                        (OUTPUT_VALUE, json.dumps({"retriever_output": "retriever-output"})),
                        (OUTPUT_MIME_TYPE, JSON),
                    )
                ),
            ),
            {"retriever_output": "retriever-output"},
            id="retriever-span-with-json-output-and-no-retrieval-documents",
        ),
        pytest.param(
            Span(
                span_kind=CHAIN,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, json.dumps({"chain_input": "chain-input"})),
                        (INPUT_MIME_TYPE, JSON),
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                    )
                ),
            ),
            {"output": "plain-text-output"},
            id="chain-span-with-plain-text-output",
        ),
        pytest.param(
            Span(
                span_kind=CHAIN,
                attributes=unflatten(
                    (
                        (INPUT_VALUE, json.dumps({"chain_input": "chain-input"})),
                        (INPUT_MIME_TYPE, JSON),
                        (OUTPUT_VALUE, json.dumps({"chain_output": "chain-output"})),
                        (OUTPUT_MIME_TYPE, JSON),
                    )
                ),
            ),
            {"chain_output": "chain-output"},
            id="chain-span-with-json-output",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "No tools here."),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {"messages": [{"content": "No tools here.", "role": "assistant"}]},
            id="llm-span-without-tools",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "Calling get_weather."),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_NAME}",
                            "get_weather",
                        ),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            '{"location": "SF"}',
                        ),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.1.{TOOL_CALL_FUNCTION_NAME}",
                            "get_time",
                        ),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.1.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            "{}",
                        ),
                    )
                ),
            ),
            {
                "messages": [
                    {
                        "content": "Calling get_weather.",
                        "role": "assistant",
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "get_weather",
                                    "arguments": {"location": "SF"},
                                },
                            },
                            {
                                "function": {
                                    "name": "get_time",
                                    "arguments": {},
                                },
                            },
                        ],
                    }
                ]
            },
            id="llm-span-with-output-messages-and-tool-calls-including-id",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        # Responses API: first item = message with content
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENTS}.0.{MESSAGE_CONTENT_TEXT}",
                            "I will call the tool.",
                        ),
                        # Second item = function_call (single tool at MESSAGE_TOOL_CALLS.0)
                        (f"{LLM_OUTPUT_MESSAGES}.1.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.1.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_NAME}",
                            "fetch_data",
                        ),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.1.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            '{"q": "x"}',
                        ),
                        # Third item = another function_call
                        (f"{LLM_OUTPUT_MESSAGES}.2.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.2.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_NAME}",
                            "other_tool",
                        ),
                        (
                            f"{LLM_OUTPUT_MESSAGES}.2.{MESSAGE_TOOL_CALLS}.0.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            "{}",
                        ),
                    )
                ),
            ),
            {
                "messages": [
                    {
                        "content": "I will call the tool.",
                        "role": "assistant",
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "fetch_data",
                                    "arguments": {"q": "x"},
                                },
                            },
                            {
                                "function": {
                                    "name": "other_tool",
                                    "arguments": {},
                                },
                            },
                        ],
                    }
                ]
            },
            id="llm-span-responses-api-merged-message-and-tool-calls",
        ),
    ],
)
def test_get_dataset_example_output(span: Span, expected_output_value: dict[str, Any]) -> None:
    output_value = get_dataset_example_output(span)
    assert expected_output_value == output_value


@pytest.mark.parametrize(
    "span, expected_output_value",
    [
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "I can help with weather."),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{SpanAttributes.LLM_TOOLS}.0.{ToolAttributes.TOOL_JSON_SCHEMA}",
                            json.dumps(
                                {
                                    "type": "function",
                                    "function": {
                                        "name": "get_weather",
                                        "description": "Get current weather",
                                        "parameters": {
                                            "type": "object",
                                            "properties": {"location": {"type": "string"}},
                                            "required": ["location"],
                                        },
                                    },
                                }
                            ),
                        ),
                    )
                ),
            ),
            {
                "messages": [{"content": "I can help with weather.", "role": "assistant"}],
                "available_tools": [
                    {
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "description": "Get current weather",
                            "parameters": {
                                "type": "object",
                                "properties": {"location": {"type": "string"}},
                                "required": ["location"],
                            },
                        },
                    }
                ],
            },
            id="llm-span-with-output-messages-and-single-tool",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "I have multiple tools."),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                        (
                            f"{SpanAttributes.LLM_TOOLS}.0.{ToolAttributes.TOOL_JSON_SCHEMA}",
                            json.dumps(
                                {
                                    "type": "function",
                                    "function": {
                                        "name": "get_weather",
                                        "parameters": {"type": "object"},
                                    },
                                }
                            ),
                        ),
                        (
                            f"{SpanAttributes.LLM_TOOLS}.1.{ToolAttributes.TOOL_JSON_SCHEMA}",
                            json.dumps(
                                {
                                    "type": "function",
                                    "function": {
                                        "name": "send_email",
                                        "parameters": {"type": "object"},
                                    },
                                }
                            ),
                        ),
                    )
                ),
            ),
            {
                "messages": [{"content": "I have multiple tools.", "role": "assistant"}],
                "available_tools": [
                    {
                        "type": "function",
                        "function": {"name": "get_weather", "parameters": {"type": "object"}},
                    },
                    {
                        "type": "function",
                        "function": {"name": "send_email", "parameters": {"type": "object"}},
                    },
                ],
            },
            id="llm-span-with-output-messages-and-multiple-tools",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", "No tools here."),
                        (f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"),
                    )
                ),
            ),
            {
                "messages": [{"content": "No tools here.", "role": "assistant"}],
                "available_tools": [],
            },
            id="llm-span-without-tools",
        ),
        pytest.param(
            Span(
                span_kind=LLM,
                attributes=unflatten(
                    (
                        (OUTPUT_VALUE, "plain-text-output"),
                        (OUTPUT_MIME_TYPE, TEXT),
                        (
                            f"{SpanAttributes.LLM_TOOLS}.0.{ToolAttributes.TOOL_JSON_SCHEMA}",
                            json.dumps(
                                {
                                    "type": "function",
                                    "function": {
                                        "name": "calculator",
                                        "parameters": {"type": "object"},
                                    },
                                }
                            ),
                        ),
                    )
                ),
            ),
            {
                "output": "plain-text-output",
                "available_tools": [
                    {
                        "type": "function",
                        "function": {"name": "calculator", "parameters": {"type": "object"}},
                    }
                ],
            },
            id="llm-span-with-plain-text-output-and-tools",
        ),
    ],
)
def test_get_experiment_example_output(span: Span, expected_output_value: dict[str, Any]) -> None:
    output_value = get_experiment_example_output(span)
    assert expected_output_value == output_value
