change the syntax of exec

minillmflow/__init__.py

@@ -7,12 +7,12 @@ class BaseNode:
         if action in self.successors: warnings.warn(f"Overwriting successor for action '{action}'")
         self.successors[action]=node;return node
     def prep(self,shared): return None
-    def exec(self,shared,prep_res): return None
-    def _exec(self,shared,prep_res): return self.exec(shared,prep_res)
+    def exec(self,prep_res): return None
+    def _exec(self,prep_res): return self.exec(prep_res)
     def post(self,shared,prep_res,exec_res): return "default"
     def _run(self,shared):
         prep_res=self.prep(shared)
-        exec_res=self._exec(shared,prep_res)
+        exec_res=self._exec(prep_res)
         return self.post(shared,prep_res,exec_res)
     def run(self,shared):
         if self.successors: warnings.warn("Node won't run successors. Use a parent Flow instead.")
@@ -30,16 +30,16 @@ class Node(BaseNode):
     def __init__(self,max_retries=1): 
         super().__init__()
         self.max_retries=max_retries
-    def process_after_fail(self,shared,prep_res,exc): raise exc
-    def _exec(self,shared,prep_res):
+    def process_after_fail(self,prep_res,exc): raise exc
+    def _exec(self,prep_res):
         for i in range(self.max_retries):
-            try:return super()._exec(shared,prep_res)
+            try:return super()._exec(prep_res)
             except Exception as e:
-                if i==self.max_retries-1:return self.process_after_fail(shared,prep_res,e)
+                if i==self.max_retries-1:return self.process_after_fail(prep_res,e)
 
 class BatchNode(Node):
     def prep(self,shared): return []
-    def _exec(self,shared,items): return [super(Node,self)._exec(shared,i) for i in items]
+    def _exec(self,items): return [super(Node,self)._exec(i) for i in items]
 
 class Flow(BaseNode):
     def __init__(self,start):
@@ -47,23 +47,24 @@ class Flow(BaseNode):
         self.start=start
     def get_next_node(self,curr,action):
         nxt=curr.successors.get(action if action is not None else "default")
-        if not nxt and curr.successors: 
-            warnings.warn(f"Flow ends: action '{action}' not found in {list(curr.successors)}")
+        if not nxt and curr.successors: warnings.warn(f"Flow ends: action '{action}' not found in {list(curr.successors)}")
         return nxt
-    def _exec(self,shared,params=None):
+    def _orchestrate(self,shared,params=None):
         curr,p=self.start,(params if params else {**self.params})
         while curr:
             curr.set_params(p)
-            c=curr._run(shared)
-            curr=self.get_next_node(curr,c)
-    def exec(self,shared,prep_res): 
+            curr=self.get_next_node(curr,curr._run(shared))
+    def _run(self,shared):
+        self._orchestrate(shared)
+        return self.post(shared,self.prep(shared),None)
+    def exec(self,prep_res): 
         raise RuntimeError("Flow should not exec directly. Create a child Node instead.")
 
 class BatchFlow(Flow):
     def prep(self,shared): return []
     def _run(self,shared):
         prep_res=self.prep(shared)
-        for batch_params in prep_res:self._exec(shared,{**self.params,**batch_params})
+        for batch_params in prep_res:self._orchestrate(shared,{**self.params,**batch_params})
         return self.post(shared,prep_res,None)
 
 class AsyncNode(Node):
@@ -77,24 +78,23 @@ class AsyncNode(Node):
         return await self._run_async(shared)
     async def _run_async(self,shared):
         prep_res=self.prep(shared)
-        exec_res=self._exec(shared,prep_res)
+        exec_res=self._exec(prep_res)
         return await self.post_async(shared,prep_res,exec_res)
     def _run(self,shared): raise RuntimeError("AsyncNode should run using run_async instead.")
 
 class AsyncFlow(Flow,AsyncNode):
-    async def _exec_async(self,shared,params=None):
+    async def _orchestrate_async(self,shared,params=None):
         curr,p=self.start,(params if params else {**self.params})
         while curr:
             curr.set_params(p)
             c=await curr._run_async(shared) if hasattr(curr,"run_async") else curr._run(shared)
             curr=self.get_next_node(curr,c)
     async def _run_async(self,shared):
-        prep_res=self.prep(shared)
-        await self._exec_async(shared)
-        return await self.post_async(shared,prep_res,None)
+        await self._orchestrate_async(shared)
+        return await self.post_async(shared,self.prep(shared),None)
 
 class BatchAsyncFlow(BatchFlow,AsyncFlow):
     async def _run_async(self,shared):
         prep_res=self.prep(shared)
-        for batch_params in prep_res:await self._exec_async(shared,{**self.params,**batch_params})
+        for batch_params in prep_res:await self._orchestrate_async(shared,{**self.params,**batch_params})
         return await self.post_async(shared,prep_res,None)