File size: 2,533 Bytes
5e4f9b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
#!/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()) |