test_structured_output.py

#!/usr/bin/env python3
"""
Test script to verify the Pydantic structured output functionality.
"""

import asyncio
import sys
import os

# Add the current directory to the path so we can import modules
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from tools.tavily_search_tool import ReportOutput, get_structured_report_from_state
from datetime import datetime

async def test_structured_output():
    """Test the structured report output functionality."""
    
    print("🧪 Testing Pydantic ReportOutput model...")
    
    # Test creating a ReportOutput instance
    test_report = ReportOutput(
        title="Test Report",
        abstract="This is a test abstract for our report.",
        content="""# Test Report

## Introduction
This is the introduction section.

## Main Content
This is the main content section with some details.

## Conclusion
This is the conclusion section.
""",
        sections=["Introduction", "Main Content", "Conclusion"],
        word_count=25,
        sources_used=["Source 1", "Source 2"]
    )
    
    print(f"✅ Created ReportOutput successfully!")
    print(f"📋 Title: {test_report.title}")
    print(f"📝 Abstract: {test_report.abstract}")
    print(f"📊 Word Count: {test_report.word_count}")
    print(f"🗂️ Sections: {test_report.sections}")
    print(f"📚 Sources: {test_report.sources_used}")
    print(f"⏰ Generated at: {test_report.generated_at}")
    
    # Test serialization
    print("\n🔄 Testing serialization...")
    serialized = test_report.model_dump()
    print(f"✅ Serialized data keys: {list(serialized.keys())}")
    
    # Test state extraction
    print("\n🔍 Testing state extraction...")
    mock_state = {
        "structured_report": serialized,
        "other_data": "some value"
    }
    
    extracted_report = get_structured_report_from_state(mock_state)
    if extracted_report:
        print(f"✅ Successfully extracted report from state!")
        print(f"📋 Extracted title: {extracted_report.title}")
        print(f"📊 Extracted word count: {extracted_report.word_count}")
    else:
        print("❌ Failed to extract report from state")
    
    # Test with empty state
    empty_report = get_structured_report_from_state({})
    if empty_report is None:
        print("✅ Correctly returned None for empty state")
    else:
        print("❌ Should have returned None for empty state")
    
    print("\n🎉 All tests completed!")

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