tests/test_structured_output.py

"""
Test structured output parsing for agent responses.
Tests the parse_agent_response function to ensure it correctly formats
images, audio, and species information from raw agent text.
"""

import asyncio
import sys
from pathlib import Path

# Add parent directory to sys.path
parent_dir = Path(__file__).parent.parent
sys.path.insert(0, str(parent_dir))

from langgraph_agent.structured_output import parse_agent_response


async def test_structured_output():
    """Test structured output parsing with various response formats."""

    print("=" * 60)
    print("TESTING STRUCTURED OUTPUT PARSING")
    print("=" * 60)

    # Test Case 1: Response with images and species identification
    print("\n[TEST 1]: Response with images and species")
    print("-" * 40)

    test_response_1 = """
    Based on the image, I can identify this as a Northern Cardinal. The bright red plumage and distinctive crest are characteristic of this species.

    Here are some reference images:
    ![Northern Cardinal](https://images.unsplash.com/photo-1578662996442-48f60103fc96?w=400)
    ![Cardinal in snow](https://images.unsplash.com/photo-1551085254-e96b210db58a?w=400)

    The Northern Cardinal is a beautiful songbird commonly found in North America.
    """

    result_1 = await parse_agent_response(
        raw_response=test_response_1,
        provider="openai",
        api_key="test-key",
        model="gpt-4o-mini"
    )

    print("Input length:", len(test_response_1))
    print("Output length:", len(result_1))
    print("Contains markdown images:", "![Northern Cardinal]" in result_1)
    print("Contains species name:", "Northern Cardinal" in result_1)
    print("✅ Test 1 completed")

    # Test Case 2: Response with audio recordings
    print("\n[TEST 2]: Response with audio recordings")
    print("-" * 40)

    test_response_2 = """
    This appears to be an American Robin. You can listen to its distinctive song here:

    Listen to the robin: https://xeno-canto.org/12345/download

    Another recording: https://xeno-canto.org/67890/download

    The American Robin is known for its cheerful song that signals the arrival of spring.
    """

    result_2 = await parse_agent_response(
        raw_response=test_response_2,
        provider="openai",
        api_key="test-key",
        model="gpt-4o-mini"
    )

    print("Input length:", len(test_response_2))
    print("Output length:", len(result_2))
    print("Contains audio links:", "[Listen to recording" in result_2)
    print("Contains xeno-canto links:", "xeno-canto.org" in result_2)
    print("✅ Test 2 completed")

    # Test Case 3: Response with no media (should return original)
    print("\n[TEST 3]: Response with no media")
    print("-" * 40)

    test_response_3 = """
    This appears to be a Blue Jay. Blue Jays are intelligent birds known for their problem-solving abilities and distinctive calls.

    They are commonly found in North American forests and suburban areas.
    """

    result_3 = await parse_agent_response(
        raw_response=test_response_3,
        provider="openai",
        api_key="test-key",
        model="gpt-4o-mini"
    )

    print("Input length:", len(test_response_3))
    print("Output length:", len(result_3))
    print("Output matches input:", result_3.strip() == test_response_3.strip())
    print("✅ Test 3 completed")

    # Test Case 4: Response with mixed URLs (images and audio)
    print("\n[TEST 4]: Response with mixed media")
    print("-" * 40)

    test_response_4 = """
    This is definitely a Scarlet Tanager. Here are some photos and recordings:

    Photo: https://example.com/tanager.jpg
    Another photo: https://example.com/tanager2.png

    Song recording: https://xeno-canto.org/11111/download
    Call recording: https://example.com/tanager.mp3

    Scarlet Tanagers are known for their striking red plumage.
    """

    result_4 = await parse_agent_response(
        raw_response=test_response_4,
        provider="openai",
        api_key="test-key",
        model="gpt-4o-mini"
    )

    print("Input length:", len(test_response_4))
    print("Output length:", len(result_4))
    print("Contains image section:", "### Images" in result_4)
    print("Contains audio section:", "### Audio Recordings" in result_4)
    print("✅ Test 4 completed")

    print("\n" + "=" * 60)
    print("✅ ALL STRUCTURED OUTPUT TESTS PASSED")
    print("=" * 60)


if __name__ == "__main__":
    asyncio.run(test_structured_output())