tests

minillmflow/__init__.py

@@ -6,10 +6,10 @@ class BaseNode:
     #  process(): this is for the LLM call, and should be idempotent for retries
     #  postprocess(): this is to summarize the result and retrun the condition for the successor node
     def __init__(self):
-        self.parameters, self.successors = {}, {}
+        self.params, self.successors = {}, {}
 
-    def set_parameters(self, params): # make sure params is immutable
-        self.parameters = params # must be immutable during pre/post/process
+    def set_params(self, params): # make sure params is immutable
+        self.params = params # must be immutable during pre/post/process
 
     def add_successor(self, node, condition="default"):
         if condition in self.successors:
@@ -30,12 +30,12 @@ class BaseNode:
     def postprocess(self, shared_storage, prep_result, proc_result):
         return "default" # condition for next node
     
-    def _run(self, shared_storage=None):
-        prep = self.preprocess(shared_storage)
-        proc = self._process(shared_storage, prep)
-        return self.postprocess(shared_storage, prep, proc)
+    def _run(self, shared_storage):
+        prep_result = self.preprocess(shared_storage)
+        proc_result = self._process(shared_storage, prep_result)
+        return self.postprocess(shared_storage, prep_result, proc_result)
     
-    def run(self, shared_storage=None):
+    def run(self, shared_storage):
         if self.successors:
             warnings.warn("This node has successor nodes. To run its successors, wrap this node in a parent Flow and use that Flow.run() instead.")
         return self._run(shared_storage)
@@ -91,37 +91,28 @@ class BatchNode(Node):
         return results
 
 class AsyncNode(Node):
-    """
-    A Node whose postprocess step is async.
-    You can also override process() to be async if needed.
-    """
     def postprocess(self, shared_storage, prep_result, proc_result):
-        # Not used in async workflow; define postprocess_async() instead.
         raise NotImplementedError("AsyncNode requires postprocess_async, and should be run in an AsyncFlow")
     
     async def postprocess_async(self, shared_storage, prep_result, proc_result):
-        """
-        Async version of postprocess. By default, returns "default".
-        Override as needed.
-        """
         await asyncio.sleep(0)  # trivial async pause (no-op)
         return "default"
     
-    async def run_async(self, shared_storage=None):
+    async def run_async(self, shared_storage):
         if self.successors:
             warnings.warn("This node has successor nodes. To run its successors, wrap this node in a parent AsyncFlow and use that AsyncFlow.run_async() instead.")
         return await self._run_async(shared_storage)
         
-    async def _run_async(self, shared_storage=None):
-        prep = self.preprocess(shared_storage)
-        proc = self._process(shared_storage, prep)
-        return await self.postprocess_async(shared_storage, prep, proc)
+    async def _run_async(self, shared_storage):
+        prep_result = self.preprocess(shared_storage)
+        proc_result = self._process(shared_storage, prep_result)
+        return await self.postprocess_async(shared_storage, prep_result, proc_result)
     
-    def _run(self, shared_storage=None):
-        raise RuntimeError("AsyncNode requires run_async, and should be run in an AsyncFlow")
+    def _run(self, shared_storage):
+        raise RuntimeError("AsyncNode requires asynchronous execution. Use 'await node.run_async()' if inside an async function, or 'asyncio.run(node.run_async())' if in synchronous code.")
 
 class BaseFlow(BaseNode):
-    def __init__(self, start_node=None):
+    def __init__(self, start_node):
         super().__init__()
         self.start_node = start_node
 
@@ -134,28 +125,26 @@ class BaseFlow(BaseNode):
         return next_node
     
 class Flow(BaseFlow):
-    def _process_flow(self, shared_storage):
+    def _process(self, shared_storage, params=None):
         current_node = self.start_node
+        params = params if params is not None else self.params.copy()
+        
         while current_node:
-            # Pass down the Flow's parameters to the current node
-            current_node.set_parameters(self.parameters)
-            # Synchronous run
+            current_node.set_params(params)
             condition = current_node._run(shared_storage)
-            # Decide next node
             current_node = self.get_next_node(current_node, condition)
-
-    def _run(self, shared_storage=None):
-        prep_result = self.preprocess(shared_storage)
-        self._process_flow(shared_storage)
-        return self.postprocess(shared_storage, prep_result, None)
-
-class AsyncFlow(BaseFlow):
-    async def _process_flow_async(self, shared_storage):
+            
+    def process(self, shared_storage, prep_result):
+        raise NotImplementedError("Flow should not process directly")
+        
+class AsyncFlow(BaseFlow, AsyncNode):
+    async def _process_async(self, shared_storage, params=None):
         current_node = self.start_node
+        params = params if params is not None else self.params.copy()
+        
         while current_node:
-            current_node.set_parameters(self.parameters)
+            current_node.set_params(params)
 
-            # If node is async-capable, call run_async; otherwise run sync
             if hasattr(current_node, "run_async") and callable(current_node.run_async):
                 condition = await current_node._run_async(shared_storage)
             else:
@@ -163,53 +152,33 @@ class AsyncFlow(BaseFlow):
 
             current_node = self.get_next_node(current_node, condition)
 
-    async def _run_async(self, shared_storage=None):
+    async def _run_async(self, shared_storage):
         prep_result = self.preprocess(shared_storage)
-        await self._process_flow_async(shared_storage)
-        return self.postprocess(shared_storage, prep_result, None)
+        await self._process_async(shared_storage)
+        return await self.postprocess_async(shared_storage, prep_result, None)
 
-    def _run(self, shared_storage=None):
-        try:
-            return asyncio.run(self._run_async(shared_storage))
-        except RuntimeError as e:
-            raise RuntimeError("If you are running in Jupyter, please use `await run_async()` instead of `run()`.") from e
-    
 class BaseBatchFlow(BaseFlow):
     def preprocess(self, shared_storage):
-        return []
+        return [] # return an iterable of parameter dictionaries
 
 class BatchFlow(BaseBatchFlow, Flow):
-    def _run(self, shared_storage=None):
+    def _run(self, shared_storage):
         prep_result = self.preprocess(shared_storage)
-        all_results = []
-
-        # For each set of parameters (or items) we got from preprocess
+        
         for param_dict in prep_result:
-            # Merge param_dict into the Flow's parameters
-            original_params = self.parameters.copy()
-            self.parameters.update(param_dict)
-
-            # Run from the start node to end
-            self._process_flow(shared_storage)
-
-            # Optionally collect results from shared_storage or a custom method
-            all_results.append(f"Finished run with parameters: {param_dict}")
-
-            # Reset the parameters if needed
-            self.parameters = original_params
+            merged_params = self.params.copy()
+            merged_params.update(param_dict)
+            self._process(shared_storage, params=merged_params)
+        
+        return self.postprocess(shared_storage, prep_result, None)
 
 class BatchAsyncFlow(BaseBatchFlow, AsyncFlow):
-    async def _run_async(self, shared_storage=None):
+    async def _run_async(self, shared_storage):
         prep_result = self.preprocess(shared_storage)
-        all_results = []
-
+        
         for param_dict in prep_result:
-            original_params = self.parameters.copy()
-            self.parameters.update(param_dict)
-
-            await self._process_flow_async(shared_storage)
-
-            all_results.append(f"Finished async run with parameters: {param_dict}")
-
-            # Reset back to original parameters if needed
-            self.parameters = original_params
+            merged_params = self.params.copy()
+            merged_params.update(param_dict)
+            await self._process_async(shared_storage, params=merged_params)
+            
+        return await self.postprocess_async(shared_storage, prep_result, None)