File size: 3,503 Bytes
741c3da |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import asyncio
from llama_index.core.workflow import Context
from llama_index.core.agent.workflow import AgentWorkflow
from tools.tavily_search_tool import (
search_web, record_notes, write_report, review_report,
SearchWebArgs, RecordNotesArgs, WriteReportArgs, ReviewReportArgs
)
async def test_tools():
"""Test all tools with the new Pydantic signatures."""
print("Testing tools with new Pydantic model arguments...")
# Create a simple workflow for context
workflow = AgentWorkflow(agents=[], root_agent=None)
ctx = Context(workflow)
# Initialize state
await ctx.set("state", {
"research_notes": {},
"report_content": "Not written yet.",
"review": "Review required.",
})
print("\n1. Testing search_web...")
try:
search_args = SearchWebArgs(query="artificial intelligence history")
search_result = await search_web(search_args)
print(f"✅ search_web worked! Result length: {len(search_result)}")
print(f"Preview: {search_result[:200]}...")
except Exception as e:
print(f"❌ search_web failed: {e}")
print("\n2. Testing record_notes...")
try:
notes_args = RecordNotesArgs(
notes="Test research notes about AI history",
notes_title="AI History Overview"
)
notes_result = await record_notes(ctx, notes_args)
print(f"✅ record_notes worked! Result: {notes_result}")
# Check state
state = await ctx.get("state")
print(f"State after notes: {list(state.keys())}")
print(f"Research notes: {state.get('research_notes', {})}")
except Exception as e:
print(f"❌ record_notes failed: {e}")
print("\n3. Testing write_report...")
try:
report_args = WriteReportArgs(
report_content="""# Artificial Intelligence History
## Introduction
This is a test report about AI history.
## Early Development
AI began in the 1950s with researchers like Alan Turing.
## Modern Era
Today, AI includes machine learning and deep learning.
## Conclusion
AI continues to evolve rapidly.""",
title="Test AI History Report"
)
report_result = await write_report(ctx, report_args)
print(f"✅ write_report worked! Result: {report_result}")
# Check state
state = await ctx.get("state")
print(f"Report content length: {len(state.get('report_content', ''))}")
print(f"Has structured report: {'structured_report' in state}")
except Exception as e:
print(f"❌ write_report failed: {e}")
print("\n4. Testing review_report...")
try:
review_args = ReviewReportArgs(review="APPROVED: The report looks good!")
review_result = await review_report(ctx, review_args)
print(f"✅ review_report worked! Result: {review_result}")
# Check final state
state = await ctx.get("state")
print(f"Final review: {state.get('review', 'No review')}")
except Exception as e:
print(f"❌ review_report failed: {e}")
print("\n5. Final state check...")
final_state = await ctx.get("state")
print(f"Final state keys: {list(final_state.keys())}")
print(f"Research notes count: {len(final_state.get('research_notes', {}))}")
print(f"Report written: {final_state.get('report_content', 'Not written') != 'Not written yet.'}")
if __name__ == "__main__":
asyncio.run(test_tools()) |