refactor

2024-12-26 02:35:45 +00:00
parent 5f95f23bc9
commit 812c041bb4
2 changed files with 257 additions and 165 deletions
@@ -0,0 +1,176 @@
+import unittest
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent.parent))
+from minillmflow import AsyncNode, BatchAsyncFlow
+
+class AsyncDataProcessNode(AsyncNode):
+    def process(self, shared_storage, prep_result):
+        key = self.params.get('key')
+        data = shared_storage['input_data'][key]
+        if 'results' not in shared_storage:
+            shared_storage['results'] = {}
+        shared_storage['results'][key] = data
+        return data
+
+    async def postprocess_async(self, shared_storage, prep_result, proc_result):
+        await asyncio.sleep(0.01)  # Simulate async work
+        key = self.params.get('key')
+        shared_storage['results'][key] = proc_result * 2  # Double the value
+        return "processed"
+
+class AsyncErrorNode(AsyncNode):
+    async def postprocess_async(self, shared_storage, prep_result, proc_result):
+        key = self.params.get('key')
+        if key == 'error_key':
+            raise ValueError(f"Async error processing key: {key}")
+        return "processed"
+
+class TestAsyncBatchFlow(unittest.TestCase):
+    def setUp(self):
+        self.process_node = AsyncDataProcessNode()
+
+    def test_basic_async_batch_processing(self):
+        """Test basic async batch processing with multiple keys"""
+        class SimpleTestAsyncBatchFlow(BatchAsyncFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        shared_storage = {
+            'input_data': {
+                'a': 1,
+                'b': 2,
+                'c': 3
+            }
+        }
+
+        flow = SimpleTestAsyncBatchFlow(start_node=self.process_node)
+        asyncio.run(flow.run_async(shared_storage))
+
+        expected_results = {
+            'a': 2,  # 1 * 2
+            'b': 4,  # 2 * 2
+            'c': 6   # 3 * 2
+        }
+        self.assertEqual(shared_storage['results'], expected_results)
+
+    def test_empty_async_batch(self):
+        """Test async batch processing with empty input"""
+        class EmptyTestAsyncBatchFlow(BatchAsyncFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        shared_storage = {
+            'input_data': {}
+        }
+
+        flow = EmptyTestAsyncBatchFlow(start_node=self.process_node)
+        asyncio.run(flow.run_async(shared_storage))
+
+        self.assertEqual(shared_storage.get('results', {}), {})
+
+    def test_async_error_handling(self):
+        """Test error handling during async batch processing"""
+        class ErrorTestAsyncBatchFlow(BatchAsyncFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        shared_storage = {
+            'input_data': {
+                'normal_key': 1,
+                'error_key': 2,
+                'another_key': 3
+            }
+        }
+
+        flow = ErrorTestAsyncBatchFlow(start_node=AsyncErrorNode())
+        
+        with self.assertRaises(ValueError):
+            asyncio.run(flow.run_async(shared_storage))
+
+    def test_nested_async_flow(self):
+        """Test async batch processing with nested flows"""
+        class AsyncInnerNode(AsyncNode):
+            async def postprocess_async(self, shared_storage, prep_result, proc_result):
+                key = self.params.get('key')
+                if 'intermediate_results' not in shared_storage:
+                    shared_storage['intermediate_results'] = {}
+                shared_storage['intermediate_results'][key] = shared_storage['input_data'][key] + 1
+                await asyncio.sleep(0.01)
+                return "next"
+
+        class AsyncOuterNode(AsyncNode):
+            async def postprocess_async(self, shared_storage, prep_result, proc_result):
+                key = self.params.get('key')
+                if 'results' not in shared_storage:
+                    shared_storage['results'] = {}
+                shared_storage['results'][key] = shared_storage['intermediate_results'][key] * 2
+                await asyncio.sleep(0.01)
+                return "done"
+
+        class NestedAsyncBatchFlow(BatchAsyncFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        # Create inner flow
+        inner_node = AsyncInnerNode()
+        outer_node = AsyncOuterNode()
+        inner_node - "next" >> outer_node
+
+        shared_storage = {
+            'input_data': {
+                'x': 1,
+                'y': 2
+            }
+        }
+
+        flow = NestedAsyncBatchFlow(start_node=inner_node)
+        asyncio.run(flow.run_async(shared_storage))
+
+        expected_results = {
+            'x': 4,  # (1 + 1) * 2
+            'y': 6   # (2 + 1) * 2
+        }
+        self.assertEqual(shared_storage['results'], expected_results)
+
+    def test_custom_async_parameters(self):
+        """Test async batch processing with additional custom parameters"""
+        class CustomParamAsyncNode(AsyncNode):
+            async def postprocess_async(self, shared_storage, prep_result, proc_result):
+                key = self.params.get('key')
+                multiplier = self.params.get('multiplier', 1)
+                await asyncio.sleep(0.01)
+                if 'results' not in shared_storage:
+                    shared_storage['results'] = {}
+                shared_storage['results'][key] = shared_storage['input_data'][key] * multiplier
+                return "done"
+
+        class CustomParamAsyncBatchFlow(BatchAsyncFlow):
+            def preprocess(self, shared_storage):
+                return [{
+                    'key': k,
+                    'multiplier': i + 1
+                } for i, k in enumerate(shared_storage['input_data'].keys())]
+
+        shared_storage = {
+            'input_data': {
+                'a': 1,
+                'b': 2,
+                'c': 3
+            }
+        }
+
+        flow = CustomParamAsyncBatchFlow(start_node=CustomParamAsyncNode())
+        asyncio.run(flow.run_async(shared_storage))
+
+        expected_results = {
+            'a': 1 * 1,  # first item, multiplier = 1
+            'b': 2 * 2,  # second item, multiplier = 2
+            'c': 3 * 3   # third item, multiplier = 3
+        }
+        self.assertEqual(shared_storage['results'], expected_results)
+
+if __name__ == '__main__':
+    unittest.main()