|
33 | 33 | from adalflow.core.base_data_class import DataClass |
34 | 34 | from dataclasses import dataclass |
35 | 35 |
|
36 | | -from adalflow.utils.lazy_import import setup_env |
| 36 | +from adalflow.utils import setup_env |
37 | 37 |
|
38 | 38 | # Direct Anthropic API imports |
39 | 39 | import anthropic |
40 | 40 |
|
| 41 | +import logging |
| 42 | + |
| 43 | +logging.basicConfig(level=logging.DEBUG) |
41 | 44 |
|
42 | 45 | setup_env() # Comment out to prevent FileNotFoundError when .env doesn't exist |
43 | 46 |
|
@@ -812,9 +815,8 @@ async def test_streaming(): |
812 | 815 | print(f" Event: {type(event).__name__}") |
813 | 816 | print(f" Event content: {event}") |
814 | 817 |
|
815 | | - await streaming_result.wait_for_completion() |
816 | 818 | print(f" ✓ Final completion status: {streaming_result.is_complete}") |
817 | | - print(f" ✓ Final result type: {type(streaming_result.final_result)}") |
| 819 | + print(f" ✓ Final answer type: {type(streaming_result.answer)}") |
818 | 820 |
|
819 | 821 | asyncio.run(test_streaming()) |
820 | 822 |
|
@@ -851,11 +853,12 @@ async def test_streaming(): |
851 | 853 | ) |
852 | 854 |
|
853 | 855 | async def test_data_class_streaming(): |
854 | | - # Wait for completion |
855 | | - await streaming_result.wait_for_completion() |
| 856 | + # Consume stream events to completion |
| 857 | + async for event in streaming_result.stream_events(): |
| 858 | + pass # Process events if needed |
856 | 859 |
|
857 | 860 | assert streaming_result.is_complete, "Should be complete" |
858 | | - assert streaming_result.final_result is not None, "Should have final result" |
| 861 | + assert streaming_result.answer is not None, "Should have final answer" |
859 | 862 |
|
860 | 863 | print(" ✓ Streaming completed") |
861 | 864 | print(" ✓ Final result available") |
@@ -1003,14 +1006,14 @@ async def test_complete_workflow(): |
1003 | 1006 | print(f" Event: {type(event).__name__}") |
1004 | 1007 | print(f" Event content: {event}") |
1005 | 1008 |
|
1006 | | - # Verify completion |
| 1009 | + # Verify completion - stream_events() automatically handles completion |
1007 | 1010 | assert streaming_result.is_complete, "Workflow should be complete" |
1008 | | - assert streaming_result.final_result is not None, "Should have final result" |
| 1011 | + assert streaming_result.answer is not None, "Should have final answer" |
1009 | 1012 | assert len(events_received) > 0, "Should receive events" |
1010 | 1013 |
|
1011 | 1014 | print(f" ✓ Received {len(events_received)} events") |
1012 | 1015 | print(" ✓ Workflow completed successfully") |
1013 | | - print(f" ✓ Final answer: {streaming_result.final_result.answer}") |
| 1016 | + print(f" ✓ Final answer: {streaming_result.answer}") |
1014 | 1017 |
|
1015 | 1018 | asyncio.run(test_complete_workflow()) |
1016 | 1019 |
|
@@ -1042,11 +1045,16 @@ async def test_error_handling(): |
1042 | 1045 | try: |
1043 | 1046 | async for event in streaming_result.stream_events(): |
1044 | 1047 | pass |
1045 | | - # If we get here, check if error was captured |
1046 | | - if streaming_result.final_result and hasattr( |
1047 | | - streaming_result.final_result, "error" |
| 1048 | + # Check if error was captured - stream_events() handles completion |
| 1049 | + if ( |
| 1050 | + hasattr(streaming_result, "_exception") |
| 1051 | + and streaming_result._exception |
1048 | 1052 | ): |
1049 | | - print(" ✓ Error properly captured in final result") |
| 1053 | + print( |
| 1054 | + f" ✓ Error captured in streaming result: {streaming_result._exception}" |
| 1055 | + ) |
| 1056 | + elif streaming_result.answer is None: |
| 1057 | + print(" ✓ No answer found, likely failed as expected") |
1050 | 1058 | else: |
1051 | 1059 | print(" ⚠ No error found, might have succeeded unexpectedly") |
1052 | 1060 | except Exception as e: |
@@ -1096,20 +1104,21 @@ async def compare_async_methods(): |
1096 | 1104 |
|
1097 | 1105 | astream_result = runner.astream(prompt_kwargs={"input_str": "Add 2 and 3"}) |
1098 | 1106 |
|
1099 | | - await astream_result.wait_for_completion() |
| 1107 | + # Consume stream events to completion |
| 1108 | + async for event in astream_result.stream_events(): |
| 1109 | + print("THIS IS THE EVENT", event) |
| 1110 | + pass # Process events if needed |
1100 | 1111 |
|
1101 | 1112 | print("\nAsync Compare Agent Final Results:\n") |
1102 | 1113 | print(f"acall result: {acall_result}") |
1103 | | - print(f"astream result: {astream_result.final_result}") |
| 1114 | + print(f"astream result: {astream_result.answer}") |
1104 | 1115 |
|
1105 | 1116 | # Both should complete successfully |
1106 | 1117 | assert hasattr(acall_result, "answer"), "acall should have answer" |
1107 | | - assert ( |
1108 | | - astream_result.final_result is not None |
1109 | | - ), "astream should have final result" |
| 1118 | + assert astream_result.answer is not None, "astream should have final answer" |
1110 | 1119 |
|
1111 | 1120 | print(f" ✓ acall result type: {type(acall_result)}") |
1112 | | - print(f" ✓ astream result type: {type(astream_result.final_result)}") |
| 1121 | + print(f" ✓ astream result answer type: {type(astream_result.answer)}") |
1113 | 1122 |
|
1114 | 1123 | asyncio.run(compare_async_methods()) |
1115 | 1124 |
|
|
0 commit comments