async def test_final_answer_event_format(mock_streaming_agent):
    """Test that the final answer event follows the expected format."""
    messages = [Message(role="user", content="Test query")]

    # Run the agent
    stream = mock_streaming_agent.arun(messages=messages)
    output = await collect_stream_output(stream)

    # Look for final answer event
    final_answer_events = [
        line for line in output
        if 'event: agent.final_answer' in line
    ]

    assert len(final_answer_events) > 0, "Final answer event should be emitted"

    # Check the format of the final answer event
    for event in final_answer_events:
        assert 'event: agent.final_answer' in event, "Event type should be 'agent.final_answer'"
        assert 'data: ' in event, "Event should have data payload"

        # Parse the data payload
        data_part = event.split('data: ')[1] if 'data: ' in event else event
        try:
            data = json.loads(data_part)

            # Check required fields
            assert 'id' in data, "Final answer event should have an 'id'"
            assert 'object' in data and data['object'] == 'agent.final_answer', "Event object should be 'agent.final_answer'"
            assert 'generated_answer' in data, "Final answer event should have a 'generated_answer'"
            assert 'citations' in data, "Final answer event should have 'citations'"

            # Check citation fields
            for citation in data['citations']:
                assert 'id' in citation, "Citation should have an 'id'"
                assert 'object' in citation and citation['object'] == 'citation', "Citation object should be 'citation'"
                assert 'spans' in citation, "Citation should have 'spans'"
                assert 'payload' in citation, "Citation should have a 'payload'"

                # Check spans format
                for span in citation['spans']:
                    assert 'start' in span, "Span should have 'start'"
                    assert 'end' in span, "Span should have 'end'"

                # Check payload fields
                assert 'document_id' in citation['payload'], "Payload should have 'document_id'"
                assert 'text' in citation['payload'], "Payload should have 'text'"
                assert 'metadata' in citation['payload'], "Payload should have 'metadata'"

        except json.JSONDecodeError:
            pytest.fail(f"Final answer event data is not valid JSON: {data_part}")
Test that the final answer event follows the expected format.
py/tests/unit/agent/test_agent_citations_old.py