microsoft--agent-lightning/tests/adapter/test_triplet_trace_tree.py

# Copyright (c) Microsoft. All rights reserved.

import itertools
import json
from typing import Any, Dict, List, Optional

import pytest

from agentlightning.adapter.triplet import RewardMatchPolicy, TracerTraceToTriplet, TraceTree
from agentlightning.types import Span
from agentlightning.types.tracer import SpanNames
from agentlightning.utils.otel import filter_and_unflatten_attributes

_SEQ = itertools.count()


def qwen_multimodal_attrs(response_id: str) -> Dict[str, Any]:
    """Simplified attributes derived from the provided Qwen trace dump."""
    prompt_content = json.dumps(
        [
            {"type": "text", "text": "Question: How many food items are shown in the bar graph?"},
            {"type": "image_url", "image_url": {"url": "file:///root/test.png"}},
        ]
    )
    return {
        "gen_ai.request.type": "chat",
        "gen_ai.system": "OpenAI",
        "gen_ai.request.model": "Qwen/Qwen2-VL-2B-Instruct",
        "gen_ai.request.temperature": 0.0,
        "gen_ai.request.streaming": False,
        "gen_ai.request.headers": "{'X-Stainless-Raw-Response': 'true'}",
        "gen_ai.prompt.0.role": "user",
        "gen_ai.prompt.0.content": prompt_content,
        "gen_ai.response.id": response_id,
        "gen_ai.response.model": "Qwen/Qwen2-VL-2B-Instruct",
        "gen_ai.usage.total_tokens": 12,
        "gen_ai.usage.prompt_tokens": 10,
        "gen_ai.usage.completion_tokens": 2,
        "gen_ai.completion.0.content": "The bar graph shows 10 food items.",
        "gen_ai.completion.0.finish_reason": "stop",
        "gen_ai.completion.0.role": "assistant",
        # Shortened token arrays to keep the fixture readable.
        "prompt_token_ids": (151644, 8948, 198, 2610),
        "response_token_ids": (785, 3619, 4771),
    }


def gpt_multimodal_attrs(response_id: str) -> Dict[str, Any]:
    """Simplified attributes derived from the provided GPT-4o trace dump."""
    prompt_content = json.dumps(
        [
            {"type": "text", "text": "Question: How many food items are shown in the bar graph?"},
            {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,AAA..."}},
        ]
    )
    prompt_filter_results = json.dumps(
        [
            {
                "prompt_index": 1,
                "content_filter_result": {"sexual": {"filtered": False, "severity": "safe"}},
            },
            {"prompt_index": 0, "content_filter_result": {}},
        ]
    )
    completion_filter_results = json.dumps(
        {
            "hate": {"filtered": False, "severity": "safe"},
            "violence": {"filtered": False, "severity": "safe"},
        }
    )
    return {
        "gen_ai.request.type": "chat",
        "gen_ai.system": "OpenAI",
        "gen_ai.request.model": "gpt-4.1-mini",
        "gen_ai.request.temperature": 0.0,
        "gen_ai.request.streaming": False,
        "gen_ai.request.headers": "{'X-Stainless-Raw-Response': 'true'}",
        "gen_ai.openai.system_fingerprint": "fp_3dcd5944f5",
        "gen_ai.prompt.0.role": "user",
        "gen_ai.prompt.0.content": prompt_content,
        "gen_ai.prompt.prompt_filter_results": prompt_filter_results,
        "gen_ai.response.id": response_id,
        "gen_ai.response.model": "gpt-4.1-mini-2025-04-14",
        "gen_ai.usage.total_tokens": 9,
        "gen_ai.usage.prompt_tokens": 7,
        "gen_ai.usage.completion_tokens": 2,
        "gen_ai.completion.0.finish_reason": "stop",
        "gen_ai.completion.0.role": "assistant",
        "gen_ai.completion.0.content": "The bar graph shows 13 food items.",
        "gen_ai.completion.0.content_filter_results": completion_filter_results,
    }


def make_span(
    span_id: str,
    name: str,
    *,
    parent_id: Optional[str],
    start_time: float,
    end_time: float,
    attributes: Optional[Dict[str, Any]] = None,
) -> Span:
    return Span.from_attributes(
        rollout_id="rollout-1",
        attempt_id="attempt-1",
        sequence_id=next(_SEQ),
        trace_id="trace-1",
        span_id=span_id,
        parent_id=parent_id,
        name=name,
        attributes=attributes or {},
        start_time=start_time,
        end_time=end_time,
    )


def make_llm_span(
    span_id: str,
    *,
    parent_id: str,
    start: float,
    end: float,
    prompt_ids: Optional[List[int]] = None,
    response_ids: Optional[List[int]] = None,
    response_id: Optional[str] = None,
    extra_attrs: Optional[Dict[str, Any]] = None,
) -> Span:
    attrs: Dict[str, Any] = {
        "prompt_token_ids": prompt_ids or [],
        "response_token_ids": response_ids or [],
    }
    if response_id is not None:
        attrs["gen_ai.response.id"] = response_id
    if extra_attrs:
        attrs.update(extra_attrs)
    return make_span(
        span_id,
        "openai.chat.completion",
        parent_id=parent_id,
        start_time=start,
        end_time=end,
        attributes=attrs,
    )


def reward_attributes(value: float) -> Dict[str, Any]:
    return {
        "agentops.task.output": json.dumps({"type": "reward", "value": value}),
    }


def make_trace_tree_root() -> TraceTree:
    """Create a minimal trace tree root for helper-only tests."""
    root_span = make_span(
        span_id="trace-root",
        name="agent.session",
        parent_id=None,
        start_time=0.0,
        end_time=1.0,
    )
    return TraceTree(root_span.span_id, root_span)


def test_trace_tree_from_spans_orders_children_and_agent_names():
    root = make_span(
        "root",
        "agent.session",
        parent_id=None,
        start_time=0.0,
        end_time=10.0,
        attributes={"agent.name": "primary-agent"},
    )
    llm = make_llm_span(
        "llm",
        parent_id="root",
        start=1.0,
        end=2.0,
        prompt_ids=[1, 2],
        response_ids=[3, 4],
        response_id="resp-1",
    )

    tree = TraceTree.from_spans([llm, root])

    assert tree.id == "root"
    assert [child.id for child in tree.children] == ["llm"]
    assert tree.find_id("llm") is tree.children[0]
    assert tree.names_tuple() == ("agent.session [primary-agent]", [("openai.chat.completion", [])])
    as_json = tree.to_json()
    assert as_json["children"][0]["span"]["name"] == "openai.chat.completion"


def test_trace_tree_virtual_root_for_multiple_roots():
    first_root = make_span("root-a", "agent.first", parent_id=None, start_time=0.0, end_time=5.0)
    second_root = make_span("root-b", "agent.second", parent_id=None, start_time=5.0, end_time=9.0)

    tree = TraceTree.from_spans([first_root, second_root])

    assert tree.id == "virtual-root"
    assert tree.span.name == "virtual-root"
    assert tree.start_time == first_root.start_time
    assert tree.end_time == second_root.end_time
    assert {child.id for child in tree.children} == {"root-a", "root-b"}


def test_trace_tree_handles_missing_parent_and_empty_input():
    with pytest.raises(ValueError):
        TraceTree.from_spans([])

    orphan_child = make_span(
        "child",
        "agent.child",
        parent_id="ghost-parent",
        start_time=2.0,
        end_time=4.0,
        attributes={"agent.name": "nested"},
    )
    llm = make_llm_span(
        "grandchild",
        parent_id="child",
        start=3.0,
        end=3.5,
        prompt_ids=[1],
        response_ids=[2],
        response_id="resp-nested",
    )

    tree = TraceTree.from_spans([llm, orphan_child])

    assert tree.id == "ghost-parent"
    assert tree.span.name == SpanNames.VIRTUAL.value
    assert tree.span.rollout_id == orphan_child.rollout_id
    assert [child.id for child in tree.children] == ["child"]
    assert tree.children[0].children[0].id == "grandchild"


def test_trace_tree_repair_hierarchy_moves_llm_span_under_agent():
    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
    agent = make_span(
        "agent",
        "agent.node",
        parent_id="root",
        start_time=1.0,
        end_time=9.0,
        attributes={"agent.name": "planner"},
    )
    llm = make_llm_span(
        "llm",
        parent_id="root",
        start=2.0,
        end=3.0,
        prompt_ids=[42],
        response_ids=[7],
        response_id="resp-planner",
    )

    tree = TraceTree.from_spans([root, agent, llm])
    assert any(child.id == "llm" for child in tree.children)

    tree.repair_hierarchy()

    assert not any(child.id == "llm" for child in tree.children)
    agent_node = tree.find_id("agent")
    assert agent_node is not None
    assert [child.id for child in agent_node.children] == ["llm"]


def test_trace_tree_to_trajectory_skips_empty_and_dedupes_llm_calls():
    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
    agent = make_span(
        "agent",
        "agent.node",
        parent_id="root",
        start_time=1.0,
        end_time=9.0,
        attributes={"agent.name": "primary-agent"},
    )
    first = make_llm_span(
        "llm-1",
        parent_id="agent",
        start=2.0,
        end=3.0,
        prompt_ids=[1, 2],
        response_ids=[3, 4],
        response_id="resp-1",
    )
    duplicate = make_llm_span(
        "llm-2",
        parent_id="agent",
        start=3.2,
        end=3.8,
        prompt_ids=[9],
        response_ids=[8],
        response_id="resp-1",
    )
    empty_tokens = make_llm_span(
        "llm-3",
        parent_id="agent",
        start=4.0,
        end=5.0,
        prompt_ids=[],
        response_ids=[],
        response_id="resp-2",
    )
    reward = make_span(
        "reward",
        "agent.reward",
        parent_id="agent",
        start_time=6.0,
        end_time=6.1,
        attributes=reward_attributes(0.5),
    )

    tree = TraceTree.from_spans([root, agent, first, duplicate, empty_tokens, reward])

    trajectory = tree.to_trajectory(
        agent_match="primary-agent",
        dedup_llm_call=True,
        _skip_empty_token_spans=True,
    )
    assert len(trajectory) == 1
    triplet = trajectory[0]
    assert triplet.prompt["token_ids"] == [1, 2]
    assert triplet.response["token_ids"] == [3, 4]
    assert triplet.metadata["response_id"] == "resp-1"
    assert triplet.metadata["agent_name"] == "primary-agent"
    assert triplet.reward == 0.5

    with_final_reward = tree.to_trajectory(
        agent_match="primary-agent",
        dedup_llm_call=True,
        _skip_empty_token_spans=True,
        final_reward=1.0,
    )
    assert len(with_final_reward) == 1
    assert with_final_reward[0].reward == 1.0


def test_tracer_trace_to_triplet_repair_required_for_agent_filter():
    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
    agent = make_span(
        "agent",
        "agent.node",
        parent_id="root",
        start_time=1.0,
        end_time=9.0,
        attributes={"agent.name": "planner"},
    )
    llm_outside_agent = make_llm_span(
        "llm",
        parent_id="root",
        start=2.0,
        end=3.0,
        prompt_ids=[7],
        response_ids=[8],
        response_id="resp-planner",
    )
    reward = make_span(
        "reward",
        "agent.reward",
        parent_id="agent",
        start_time=4.0,
        end_time=4.5,
        attributes=reward_attributes(0.3),
    )
    spans = [root, agent, llm_outside_agent, reward]

    adapter = TracerTraceToTriplet(agent_match="planner")
    triplets = adapter.adapt(spans)
    assert len(triplets) == 1
    assert triplets[0].metadata["agent_name"] == "planner"
    assert triplets[0].reward == 0.3

    adapter_without_repair = TracerTraceToTriplet(repair_hierarchy=False, agent_match="planner")
    assert adapter_without_repair.adapt(spans) == []


def test_tracer_trace_to_triplet_dedup_and_skip_empty_token_spans():
    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
    agent = make_span(
        "agent",
        "agent.node",
        parent_id="root",
        start_time=1.0,
        end_time=9.0,
        attributes={"agent.name": "collector"},
    )
    kept_llm = make_llm_span(
        "llm-1",
        parent_id="agent",
        start=2.0,
        end=3.0,
        prompt_ids=[10],
        response_ids=[20],
        response_id="resp-shared",
    )
    duplicate_llm = make_llm_span(
        "llm-2",
        parent_id="agent",
        start=3.5,
        end=4.2,
        prompt_ids=[99],
        response_ids=[98],
        response_id="resp-shared",
    )
    missing_tokens = make_llm_span(
        "llm-3",
        parent_id="agent",
        start=5.0,
        end=5.5,
        prompt_ids=[],
        response_ids=[],
        response_id="resp-3",
    )
    reward = make_span(
        "reward",
        "agent.reward",
        parent_id="agent",
        start_time=6.0,
        end_time=6.5,
        attributes=reward_attributes(0.25),
    )
    spans = [root, agent, kept_llm, duplicate_llm, missing_tokens, reward]

    adapter = TracerTraceToTriplet(_skip_empty_token_spans=True)
    triplets = adapter.adapt(spans)

    assert len(triplets) == 1
    assert triplets[0].prompt["token_ids"] == [10]
    assert triplets[0].response["token_ids"] == [20]
    assert triplets[0].metadata["response_id"] == "resp-shared"
    assert triplets[0].reward == 0.25


def test_trace_tree_find_llm_calls_dedupes_across_agents():
    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
    agent_a = make_span(
        "agent-a",
        "agent.node",
        parent_id="root",
        start_time=0.5,
        end_time=5.0,
        attributes={"agent.name": "vision-a"},
    )
    agent_b = make_span(
        "agent-b",
        "agent.node",
        parent_id="root",
        start_time=5.1,
        end_time=9.5,
        attributes={"agent.name": "vision-b"},
    )
    shared_response_id = "chatcmpl-shared"
    llm_a = make_span(
        "llm-a",
        "openai.chat.completion",
        parent_id="agent-a",
        start_time=1.0,
        end_time=2.0,
        attributes=qwen_multimodal_attrs(shared_response_id),
    )
    llm_b = make_span(
        "llm-b",
        "openai.chat.completion",
        parent_id="agent-b",
        start_time=6.0,
        end_time=7.0,
        attributes=gpt_multimodal_attrs(shared_response_id),
    )

    tree = TraceTree.from_spans([root, agent_a, agent_b, llm_a, llm_b])
    matches = tree.find_llm_calls(
        llm_call_match=r"openai\.chat\.completion",
        agent_match=None,
        within_matching_subtree="*",
        within_reward=False,
        within_llm_call=False,
        existing_llm_call_response_ids=set(),
    )

    assert len(matches) == 1
    assert matches[0][0].id == "llm-a"


def test_tracer_trace_to_triplet_handles_multimodal_payloads():
    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=15.0)
    agent = make_span(
        "agent",
        "agent.node",
        parent_id="root",
        start_time=0.5,
        end_time=14.5,
        attributes={"agent.name": "vision-agent"},
    )
    llm_first = make_span(
        "llm-qwen",
        "openai.chat.completion",
        parent_id="agent",
        start_time=1.0,
        end_time=2.0,
        attributes=qwen_multimodal_attrs("chatcmpl-qwen"),
    )
    llm_second = make_span(
        "llm-gpt",
        "openai.chat.completion",
        parent_id="agent",
        start_time=10.0,
        end_time=11.0,
        attributes=gpt_multimodal_attrs("chatcmpl-gpt"),
    )
    reward = make_span(
        "reward",
        "agent.reward",
        parent_id="agent",
        start_time=12.0,
        end_time=12.5,
        attributes=reward_attributes(0.7),
    )

    assert llm_first.attributes["gen_ai.request.headers"] == "{'X-Stainless-Raw-Response': 'true'}"
    assert llm_second.attributes["gen_ai.openai.system_fingerprint"] == "fp_3dcd5944f5"
    assert "gen_ai.prompt.prompt_filter_results" in llm_second.attributes
    qwen_prompt = json.loads(llm_first.attributes["gen_ai.prompt.0.content"])  # type: ignore
    assert qwen_prompt[0]["type"] == "text"
    assert qwen_prompt[1]["image_url"]["url"].startswith("file://")
    gpt_prompt = json.loads(llm_second.attributes["gen_ai.prompt.0.content"])  # type: ignore
    assert gpt_prompt[1]["image_url"]["url"].startswith("data:image/jpeg")
    assert llm_first.attributes["gen_ai.completion.0.content"] == "The bar graph shows 10 food items."
    assert llm_second.attributes["gen_ai.completion.0.content"] == "The bar graph shows 13 food items."

    adapter = TracerTraceToTriplet(agent_match="vision-agent")
    triplets = adapter.adapt([root, agent, llm_first, llm_second, reward])

    assert len(triplets) == 2
    first, second = triplets
    assert list(first.prompt["token_ids"]) == [151644, 8948, 198, 2610]
    assert list(first.response["token_ids"]) == [785, 3619, 4771]
    assert first.metadata["response_id"] == "chatcmpl-qwen"
    assert first.metadata["agent_name"] == "vision-agent"
    assert first.reward is None

    assert second.prompt["token_ids"] == []
    assert second.response["token_ids"] == []
    assert second.metadata["response_id"] == "chatcmpl-gpt"
    assert triplets[0].metadata["agent_name"] == "vision-agent"
    assert triplets[1].metadata["agent_name"] == "vision-agent"
    qwen_prompt_raw = triplets[0].prompt["raw_content"]
    assert qwen_prompt_raw == filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.prompt")
    assert triplets[0].prompt["image_urls"] == ["file:///root/test.png"]
    qwen_content = json.loads(qwen_prompt_raw[0]["content"])
    assert qwen_content[1]["image_url"]["url"] == "file:///root/test.png"
    qwen_request = filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.request")
    qwen_response = filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.response")
    assert triplets[0].metadata["request"] == qwen_request
    assert triplets[0].metadata["response"] == qwen_response
    qwen_completion = filter_and_unflatten_attributes(llm_first.attributes, "gen_ai.completion")
    assert triplets[0].response["raw_content"] == qwen_completion

    gpt_prompt_raw = triplets[1].prompt["raw_content"]
    assert gpt_prompt_raw == filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.prompt")
    gpt_content = json.loads(gpt_prompt_raw["0"]["content"])
    assert gpt_content[1]["image_url"]["url"].startswith("data:image/jpeg")
    assert triplets[1].prompt["image_urls"] == ["data:image/jpeg;base64,AAA..."]
    gpt_request = filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.request")
    gpt_response = filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.response")
    assert triplets[1].metadata["request"] == gpt_request
    assert triplets[1].metadata["response"] == gpt_response
    gpt_completion = filter_and_unflatten_attributes(llm_second.attributes, "gen_ai.completion")
    assert triplets[1].response["raw_content"] == gpt_completion
    assert second.reward == 0.7


def test_extract_prompt_image_urls_from_list_payload():
    tree = make_trace_tree_root()
    prompt_raw_content = [
        {
            "role": "user",
            "content": json.dumps(
                [
                    {"type": "text", "text": "describe the image"},
                    {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},
                ]
            ),
        },
        {
            "role": "assistant",
            "content": json.dumps(
                [
                    {"type": "text", "text": "another prompt"},
                    {"type": "image_url", "image_url": {"url": "https://example.com/b.png"}},
                ]
            ),
        },
    ]

    image_urls = tree.extract_prompt_image_urls(prompt_raw_content)

    assert image_urls == ["https://example.com/a.png", "https://example.com/b.png"]


def test_tracer_trace_to_triplet_reward_match_first_sibling():
    root = make_span("root", "session", parent_id=None, start_time=0.0, end_time=10.0)
    agent = make_span(
        "agent",
        "agent.node",
        parent_id="root",
        start_time=1.0,
        end_time=9.0,
        attributes={"agent.name": "sibling-agent"},
    )
    other_agent = make_span(
        "agent-2",
        "agent.node",
        parent_id="root",
        start_time=1.0,
        end_time=9.0,
        attributes={"agent.name": "sibling-agent"},
    )
    llm_1 = make_llm_span(
        "llm-1",
        parent_id="agent",
        start=2.0,
        end=3.0,
        prompt_ids=[1],
        response_ids=[2],
        response_id="resp-1",
    )
    reward = make_span(
        "reward",
        "agent.reward",
        parent_id="agent",
        start_time=3.5,
        end_time=3.6,
        attributes=reward_attributes(0.8),
    )
    llm_2 = make_llm_span(
        "llm-2",
        parent_id="agent-2",
        start=3.1,
        end=3.2,
        prompt_ids=[3],
        response_ids=[4],
        response_id="resp-2",
    )

    spans = [root, agent, other_agent, llm_1, reward, llm_2]

    adapter = TracerTraceToTriplet(
        agent_match="sibling-agent",
        reward_match=RewardMatchPolicy.FIRST_SIBLING,
        _skip_empty_token_spans=True,
    )
    triplets = adapter.adapt(spans)

    assert len(triplets) == 2
    t1, t2 = triplets

    assert t1.metadata["response_id"] == "resp-1"
    assert t1.reward == 0.8

    assert t2.metadata["response_id"] == "resp-2"
    assert t2.reward is None


def test_extract_prompt_image_urls_handles_numeric_dict_keys():
    tree = make_trace_tree_root()
    prompt_raw_content = {
        "1": {
            "role": "assistant",
            "content": [
                {"type": "text", "text": "second"},
                {"type": "image_url", "image_url": {"url": "https://example.com/second.png"}},
            ],
        },
        "0": {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": "https://example.com/first.png"}},
                {"type": "text", "text": "first"},
            ],
        },
    }

    image_urls = tree.extract_prompt_image_urls(prompt_raw_content)

    assert image_urls == ["https://example.com/first.png", "https://example.com/second.png"]


def test_extract_prompt_image_urls_gracefully_handles_invalid_payloads():
    tree = make_trace_tree_root()
    invalid_prompt_content = [
        {"role": "user", "content": "not-a-json"},
        {"role": "assistant", "content": json.dumps([{"type": "text", "text": "no images here"}])},
        {"role": "system"},
    ]

    assert tree.extract_prompt_image_urls(invalid_prompt_content) == []
    assert tree.extract_prompt_image_urls("unexpected-string") == []