diff --git a/.cursor/rules/core_abstraction/batch.mdc b/.cursor/rules/core_abstraction/batch.mdc
index e15d0d4d..ffbe5ffd 100644
--- a/.cursor/rules/core_abstraction/batch.mdc
+++ b/.cursor/rules/core_abstraction/batch.mdc
@@ -29,13 +29,13 @@ class MapSummaries(BatchNode):
         chunks = [content[i:i+chunk_size] for i in range(0, len(content), chunk_size)]
         return chunks
 
-    def exec(self, chunk):
-        prompt = f"Summarize this chunk in 10 words: {chunk}"
+    def exec(self, prep_res):
+        prompt = f"Summarize this chunk in 10 words: {prep_res}"
         summary = call_llm(prompt)
         return summary
 
-    def post(self, shared, prep_res, exec_res_list):
-        combined = "\n".join(exec_res_list)
+    def post(self, shared, prep_res, exec_res):
+        combined = "\n".join(exec_res)
         shared["summary"] = combined
         return "default"
 
@@ -75,8 +75,8 @@ class LoadFile(Node):
         filename = self.params["filename"]  # Important! Use self.params, not shared
         return filename
         
-    def exec(self, filename):
-        with open(filename, 'r') as f:
+    def exec(self, prep_res):
+        with open(prep_res, 'r') as f:
             return f.read()
             
     def post(self, shared, prep_res, exec_res):
@@ -89,8 +89,8 @@ class Summarize(Node):
     def prep(self, shared):
         return shared["current_file_content"]
         
-    def exec(self, content):
-        prompt = f"Summarize this file in 50 words: {content}"
+    def exec(self, prep_res):
+        prompt = f"Summarize this file in 50 words: {prep_res}"
         return call_llm(prompt)
         
     def post(self, shared, prep_res, exec_res):
@@ -154,9 +154,9 @@ class ProcessFile(Node):
         full_path = os.path.join(directory, filename)
         return full_path
         
-    def exec(self, full_path):
+    def exec(self, prep_res):
         # Process the file...
-        return f"Processed {full_path}"
+        return f"Processed {prep_res}"
         
     def post(self, shared, prep_res, exec_res):
         # Store results, perhaps indexed by path
diff --git a/.cursor/rules/core_abstraction/parallel.mdc b/.cursor/rules/core_abstraction/parallel.mdc
index 88a0b424..1783bbf4 100644
--- a/.cursor/rules/core_abstraction/parallel.mdc
+++ b/.cursor/rules/core_abstraction/parallel.mdc
@@ -27,8 +27,8 @@ class ParallelSummaries(AsyncParallelBatchNode):
         # e.g., multiple texts
         return shared["texts"]
 
-    async def exec_async(self, text):
-        prompt = f"Summarize: {text}"
+    async def exec_async(self, prep_res):
+        prompt = f"Summarize: {prep_res}"
         return await call_llm_async(prompt)
 
     async def post_async(self, shared, prep_res, exec_res_list):
diff --git a/.cursor/rules/design_pattern/agent.mdc b/.cursor/rules/design_pattern/agent.mdc
index 6c719b94..d93f6aee 100644
--- a/.cursor/rules/design_pattern/agent.mdc
+++ b/.cursor/rules/design_pattern/agent.mdc
@@ -76,8 +76,8 @@ class DecideAction(Node):
         query = shared["query"]
         return query, context
         
-    def exec(self, inputs):
-        query, context = inputs
+    def exec(self, prep_res):
+        query, context = prep_res
         prompt = f"""
 Given input: {query}
 Previous search results: {context}
@@ -110,8 +110,8 @@ class SearchWeb(Node):
     def prep(self, shared):
         return shared["search_term"]
         
-    def exec(self, search_term):
-        return search_web(search_term)
+    def exec(self, prep_res):
+        return search_web(prep_res)
     
     def post(self, shared, prep_res, exec_res):
         prev_searches = shared.get("context", [])
@@ -124,8 +124,8 @@ class DirectAnswer(Node):
     def prep(self, shared):
         return shared["query"], shared.get("context", "")
         
-    def exec(self, inputs):
-        query, context = inputs
+    def exec(self, prep_res):
+        query, context = prep_res
         return call_llm(f"Context: {context}\nAnswer: {query}")
 
     def post(self, shared, prep_res, exec_res):
diff --git a/.cursor/rules/design_pattern/mapreduce.mdc b/.cursor/rules/design_pattern/mapreduce.mdc
index 8b04f268..92470449 100644
--- a/.cursor/rules/design_pattern/mapreduce.mdc
+++ b/.cursor/rules/design_pattern/mapreduce.mdc
@@ -23,8 +23,8 @@ class SummarizeAllFiles(BatchNode):
         files_dict = shared["files"]  # e.g. 10 files
         return list(files_dict.items())  # [("file1.txt", "aaa..."), ("file2.txt", "bbb..."), ...]
 
-    def exec(self, one_file):
-        filename, file_content = one_file
+    def exec(self, prep_res):
+        filename, file_content = prep_res
         summary_text = call_llm(f"Summarize the following file:\n{file_content}")
         return (filename, summary_text)
 
@@ -35,17 +35,17 @@ class CombineSummaries(Node):
     def prep(self, shared):
         return shared["file_summaries"]
 
-    def exec(self, file_summaries):
+    def exec(self, prep_res):
         # format as: "File1: summary\nFile2: summary...\n"
         text_list = []
-        for fname, summ in file_summaries.items():
+        for fname, summ in prep_res.items():
             text_list.append(f"{fname} summary:\n{summ}\n")
         big_text = "\n---\n".join(text_list)
 
         return call_llm(f"Combine these file summaries into one final summary:\n{big_text}")
 
-    def post(self, shared, prep_res, final_summary):
-        shared["all_files_summary"] = final_summary
+    def post(self, shared, prep_res, exec_res):
+        shared["all_files_summary"] = exec_res
 
 batch_node = SummarizeAllFiles()
 combine_node = CombineSummaries()
diff --git a/.cursor/rules/design_pattern/multi_agent.mdc b/.cursor/rules/design_pattern/multi_agent.mdc
index ca3aaa8e..0aaa0d37 100644
--- a/.cursor/rules/design_pattern/multi_agent.mdc
+++ b/.cursor/rules/design_pattern/multi_agent.mdc
@@ -81,10 +81,10 @@ class AsyncHinter(AsyncNode):
             return None
         return shared["target_word"], shared["forbidden_words"], shared.get("past_guesses", [])
 
-    async def exec_async(self, inputs):
-        if inputs is None:
+    async def exec_async(self, prep_res):
+        if prep_res is None:
             return None
-        target, forbidden, past_guesses = inputs
+        target, forbidden, past_guesses = prep_res
         prompt = f"Generate hint for '{target}'\nForbidden words: {forbidden}"
         if past_guesses:
             prompt += f"\nPrevious wrong guesses: {past_guesses}\nMake hint more specific."
@@ -105,8 +105,8 @@ class AsyncGuesser(AsyncNode):
         hint = await shared["guesser_queue"].get()
         return hint, shared.get("past_guesses", [])
 
-    async def exec_async(self, inputs):
-        hint, past_guesses = inputs
+    async def exec_async(self, prep_res):
+        hint, past_guesses = prep_res
         prompt = f"Given hint: {hint}, past wrong guesses: {past_guesses}, make a new guess. Directly reply a single word:"
         guess = call_llm(prompt)
         print(f"Guesser: I guess it's - {guess}")
diff --git a/.cursor/rules/design_pattern/rag.mdc b/.cursor/rules/design_pattern/rag.mdc
index 67cf30ea..bf0f437b 100644
--- a/.cursor/rules/design_pattern/rag.mdc
+++ b/.cursor/rules/design_pattern/rag.mdc
@@ -26,9 +26,9 @@ class ChunkDocs(BatchNode):
         # A list of file paths in shared["files"]. We process each file.
         return shared["files"]
 
-    def exec(self, filepath):
+    def exec(self, prep_res):
         # read file content. In real usage, do error handling.
-        with open(filepath, "r", encoding="utf-8") as f:
+        with open(prep_res, "r", encoding="utf-8") as f:
             text = f.read()
         # chunk by 100 chars each
         chunks = []
@@ -49,8 +49,8 @@ class EmbedDocs(BatchNode):
     def prep(self, shared):
         return shared["all_chunks"]
 
-    def exec(self, chunk):
-        return get_embedding(chunk)
+    def exec(self, prep_res):
+        return get_embedding(prep_res)
 
     def post(self, shared, prep_res, exec_res_list):
         # Store the list of embeddings.
@@ -62,13 +62,13 @@ class StoreIndex(Node):
         # We'll read all embeds from shared.
         return shared["all_embeds"]
 
-    def exec(self, all_embeds):
+    def exec(self, prep_res):
         # Create a vector index (faiss or other DB in real usage).
-        index = create_index(all_embeds)
+        index = create_index(prep_res)
         return index
 
-    def post(self, shared, prep_res, index):
-        shared["index"] = index
+    def post(self, shared, prep_res, exec_res):
+        shared["index"] = exec_res
 
 # Wire them in sequence
 chunk_node = ChunkDocs()
@@ -102,40 +102,40 @@ class EmbedQuery(Node):
     def prep(self, shared):
         return shared["question"]
 
-    def exec(self, question):
-        return get_embedding(question)
+    def exec(self, prep_res):
+        return get_embedding(prep_res)
 
-    def post(self, shared, prep_res, q_emb):
-        shared["q_emb"] = q_emb
+    def post(self, shared, prep_res, exec_res):
+        shared["q_emb"] = exec_res
 
 class RetrieveDocs(Node):
     def prep(self, shared):
         # We'll need the query embedding, plus the offline index/chunks
         return shared["q_emb"], shared["index"], shared["all_chunks"]
 
-    def exec(self, inputs):
-        q_emb, index, chunks = inputs
+    def exec(self, prep_res):
+        q_emb, index, chunks = prep_res
         I, D = search_index(index, q_emb, top_k=1)
         best_id = I[0][0]
         relevant_chunk = chunks[best_id]
         return relevant_chunk
 
-    def post(self, shared, prep_res, relevant_chunk):
-        shared["retrieved_chunk"] = relevant_chunk
-        print("Retrieved chunk:", relevant_chunk[:60], "...")
+    def post(self, shared, prep_res, exec_res):
+        shared["retrieved_chunk"] = exec_res
+        print("Retrieved chunk:", exec_res[:60], "...")
 
 class GenerateAnswer(Node):
     def prep(self, shared):
         return shared["question"], shared["retrieved_chunk"]
 
-    def exec(self, inputs):
-        question, chunk = inputs
+    def exec(self, prep_res):
+        question, chunk = prep_res
         prompt = f"Question: {question}\nContext: {chunk}\nAnswer:"
         return call_llm(prompt)
 
-    def post(self, shared, prep_res, answer):
-        shared["answer"] = answer
-        print("Answer:", answer)
+    def post(self, shared, prep_res, exec_res):
+        shared["answer"] = exec_res
+        print("Answer:", exec_res)
 
 embed_qnode = EmbedQuery()
 retrieve_node = RetrieveDocs()
diff --git a/.cursor/rules/design_pattern/workflow.mdc b/.cursor/rules/design_pattern/workflow.mdc
index 98fe5dec..540d6bb1 100644
--- a/.cursor/rules/design_pattern/workflow.mdc
+++ b/.cursor/rules/design_pattern/workflow.mdc
@@ -20,17 +20,17 @@ Many real-world tasks are too complex for one LLM call. The solution is to **Tas
 ```python
 class GenerateOutline(Node):
     def prep(self, shared): return shared["topic"]
-    def exec(self, topic): return call_llm(f"Create a detailed outline for an article about {topic}")
+    def exec(self, prep_res): return call_llm(f"Create a detailed outline for an article about {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["outline"] = exec_res
 
 class WriteSection(Node):
     def prep(self, shared): return shared["outline"]
-    def exec(self, outline): return call_llm(f"Write content based on this outline: {outline}")
+    def exec(self, prep_res): return call_llm(f"Write content based on this outline: {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["draft"] = exec_res
 
 class ReviewAndRefine(Node):
     def prep(self, shared): return shared["draft"]
-    def exec(self, draft): return call_llm(f"Review and improve this draft: {draft}")
+    def exec(self, prep_res): return call_llm(f"Review and improve this draft: {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["final_article"] = exec_res
 
 # Connect nodes
diff --git a/.cursor/rules/guide_for_pocketflow.mdc b/.cursor/rules/guide_for_pocketflow.mdc
index 809e1d77..a6a90994 100644
--- a/.cursor/rules/guide_for_pocketflow.mdc
+++ b/.cursor/rules/guide_for_pocketflow.mdc
@@ -172,7 +172,7 @@ my_project/
   from utils.call_llm import call_llm
 
   class GetQuestionNode(Node):
-      def exec(self, _):
+      def exec(self, prep_res):
           # Get question directly from user input
           user_question = input("Enter your question: ")
           return user_question
@@ -187,9 +187,9 @@ my_project/
           # Read question from shared
           return shared["question"]
       
-      def exec(self, question):
+      def exec(self, prep_res):
           # Call LLM to get the answer
-          return call_llm(question)
+          return call_llm(prep_res)
       
       def post(self, shared, prep_res, exec_res):
           # Store the answer in shared
diff --git a/.cursor/rules/utility_function/llm.mdc b/.cursor/rules/utility_function/llm.mdc
index 5a1ee20c..ab9e575c 100644
--- a/.cursor/rules/utility_function/llm.mdc
+++ b/.cursor/rules/utility_function/llm.mdc
@@ -139,8 +139,8 @@ def call_llm(prompt, use_cache):
     return cached_call.__wrapped__(prompt)
 
 class SummarizeNode(Node):
-    def exec(self, text):
-        return call_llm(f"Summarize: {text}", self.cur_retry==0)
+    def exec(self, prep_res):
+        return call_llm(f"Summarize: {prep_res}", self.cur_retry==0)
 ```
 
 - Enable logging:
diff --git a/.cursorrules b/.cursorrules
index c042f3cf..ae32a68e 100644
--- a/.cursorrules
+++ b/.cursorrules
@@ -167,7 +167,7 @@ my_project/
   from utils.call_llm import call_llm
 
   class GetQuestionNode(Node):
-      def exec(self, _):
+      def exec(self, prep_res):
           # Get question directly from user input
           user_question = input("Enter your question: ")
           return user_question
@@ -182,9 +182,9 @@ my_project/
           # Read question from shared
           return shared["question"]
       
-      def exec(self, question):
+      def exec(self, prep_res):
           # Call LLM to get the answer
-          return call_llm(question)
+          return call_llm(prep_res)
       
       def post(self, shared, prep_res, exec_res):
           # Store the answer in shared
@@ -398,13 +398,13 @@ class MapSummaries(BatchNode):
         chunks = [content[i:i+chunk_size] for i in range(0, len(content), chunk_size)]
         return chunks
 
-    def exec(self, chunk):
-        prompt = f"Summarize this chunk in 10 words: {chunk}"
+    def exec(self, prep_res):
+        prompt = f"Summarize this chunk in 10 words: {prep_res}"
         summary = call_llm(prompt)
         return summary
 
-    def post(self, shared, prep_res, exec_res_list):
-        combined = "\n".join(exec_res_list)
+    def post(self, shared, prep_res, exec_res):
+        combined = "\n".join(exec_res)
         shared["summary"] = combined
         return "default"
 
@@ -1039,8 +1039,8 @@ class DecideAction(Node):
         query = shared["query"]
         return query, context
         
-    def exec(self, inputs):
-        query, context = inputs
+    def exec(self, prep_res):
+        query, context = prep_res
         prompt = f"""
 Given input: {query}
 Previous search results: {context}
@@ -1073,8 +1073,8 @@ class SearchWeb(Node):
     def prep(self, shared):
         return shared["search_term"]
         
-    def exec(self, search_term):
-        return search_web(search_term)
+    def exec(self, prep_res):
+        return search_web(prep_res)
     
     def post(self, shared, prep_res, exec_res):
         prev_searches = shared.get("context", [])
@@ -1087,8 +1087,8 @@ class DirectAnswer(Node):
     def prep(self, shared):
         return shared["query"], shared.get("context", "")
         
-    def exec(self, inputs):
-        query, context = inputs
+    def exec(self, prep_res):
+        query, context = prep_res
         return call_llm(f"Context: {context}\nAnswer: {query}")
 
     def post(self, shared, prep_res, exec_res):
@@ -1140,8 +1140,8 @@ class SummarizeAllFiles(BatchNode):
         files_dict = shared["files"]  # e.g. 10 files
         return list(files_dict.items())  # [("file1.txt", "aaa..."), ("file2.txt", "bbb..."), ...]
 
-    def exec(self, one_file):
-        filename, file_content = one_file
+    def exec(self, prep_res):
+        filename, file_content = prep_res
         summary_text = call_llm(f"Summarize the following file:\n{file_content}")
         return (filename, summary_text)
 
@@ -1152,17 +1152,17 @@ class CombineSummaries(Node):
     def prep(self, shared):
         return shared["file_summaries"]
 
-    def exec(self, file_summaries):
+    def exec(self, prep_res):
         # format as: "File1: summary\nFile2: summary...\n"
         text_list = []
-        for fname, summ in file_summaries.items():
+        for fname, summ in prep_res.items():
             text_list.append(f"{fname} summary:\n{summ}\n")
         big_text = "\n---\n".join(text_list)
 
         return call_llm(f"Combine these file summaries into one final summary:\n{big_text}")
 
-    def post(self, shared, prep_res, final_summary):
-        shared["all_files_summary"] = final_summary
+    def post(self, shared, prep_res, exec_res):
+        shared["all_files_summary"] = exec_res
 
 batch_node = SummarizeAllFiles()
 combine_node = CombineSummaries()
@@ -1217,9 +1217,9 @@ class ChunkDocs(BatchNode):
         # A list of file paths in shared["files"]. We process each file.
         return shared["files"]
 
-    def exec(self, filepath):
+    def exec(self, prep_res):
         # read file content. In real usage, do error handling.
-        with open(filepath, "r", encoding="utf-8") as f:
+        with open(prep_res, "r", encoding="utf-8") as f:
             text = f.read()
         # chunk by 100 chars each
         chunks = []
@@ -1240,8 +1240,8 @@ class EmbedDocs(BatchNode):
     def prep(self, shared):
         return shared["all_chunks"]
 
-    def exec(self, chunk):
-        return get_embedding(chunk)
+    def exec(self, prep_res):
+        return get_embedding(prep_res)
 
     def post(self, shared, prep_res, exec_res_list):
         # Store the list of embeddings.
@@ -1253,13 +1253,13 @@ class StoreIndex(Node):
         # We'll read all embeds from shared.
         return shared["all_embeds"]
 
-    def exec(self, all_embeds):
+    def exec(self, prep_res):
         # Create a vector index (faiss or other DB in real usage).
-        index = create_index(all_embeds)
+        index = create_index(prep_res)
         return index
 
-    def post(self, shared, prep_res, index):
-        shared["index"] = index
+    def post(self, shared, prep_res, exec_res):
+        shared["index"] = exec_res
 
 # Wire them in sequence
 chunk_node = ChunkDocs()
@@ -1293,40 +1293,40 @@ class EmbedQuery(Node):
     def prep(self, shared):
         return shared["question"]
 
-    def exec(self, question):
-        return get_embedding(question)
+    def exec(self, prep_res):
+        return get_embedding(prep_res)
 
-    def post(self, shared, prep_res, q_emb):
-        shared["q_emb"] = q_emb
+    def post(self, shared, prep_res, exec_res):
+        shared["q_emb"] = exec_res
 
 class RetrieveDocs(Node):
     def prep(self, shared):
         # We'll need the query embedding, plus the offline index/chunks
         return shared["q_emb"], shared["index"], shared["all_chunks"]
 
-    def exec(self, inputs):
-        q_emb, index, chunks = inputs
+    def exec(self, prep_res):
+        q_emb, index, chunks = prep_res
         I, D = search_index(index, q_emb, top_k=1)
         best_id = I[0][0]
         relevant_chunk = chunks[best_id]
         return relevant_chunk
 
-    def post(self, shared, prep_res, relevant_chunk):
-        shared["retrieved_chunk"] = relevant_chunk
-        print("Retrieved chunk:", relevant_chunk[:60], "...")
+    def post(self, shared, prep_res, exec_res):
+        shared["retrieved_chunk"] = exec_res
+        print("Retrieved chunk:", exec_res[:60], "...")
 
 class GenerateAnswer(Node):
     def prep(self, shared):
         return shared["question"], shared["retrieved_chunk"]
 
-    def exec(self, inputs):
-        question, chunk = inputs
+    def exec(self, prep_res):
+        question, chunk = prep_res
         prompt = f"Question: {question}\nContext: {chunk}\nAnswer:"
         return call_llm(prompt)
 
-    def post(self, shared, prep_res, answer):
-        shared["answer"] = answer
-        print("Answer:", answer)
+    def post(self, shared, prep_res, exec_res):
+        shared["answer"] = exec_res
+        print("Answer:", exec_res)
 
 embed_qnode = EmbedQuery()
 retrieve_node = RetrieveDocs()
@@ -1494,17 +1494,17 @@ Many real-world tasks are too complex for one LLM call. The solution is to **Tas
 ```python
 class GenerateOutline(Node):
     def prep(self, shared): return shared["topic"]
-    def exec(self, topic): return call_llm(f"Create a detailed outline for an article about {topic}")
+    def exec(self, prep_res): return call_llm(f"Create a detailed outline for an article about {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["outline"] = exec_res
 
 class WriteSection(Node):
     def prep(self, shared): return shared["outline"]
-    def exec(self, outline): return call_llm(f"Write content based on this outline: {outline}")
+    def exec(self, prep_res): return call_llm(f"Write content based on this outline: {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["draft"] = exec_res
 
 class ReviewAndRefine(Node):
     def prep(self, shared): return shared["draft"]
-    def exec(self, draft): return call_llm(f"Review and improve this draft: {draft}")
+    def exec(self, prep_res): return call_llm(f"Review and improve this draft: {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["final_article"] = exec_res
 
 # Connect nodes
@@ -1653,8 +1653,8 @@ def call_llm(prompt, use_cache):
     return cached_call.__wrapped__(prompt)
 
 class SummarizeNode(Node):
-    def exec(self, text):
-        return call_llm(f"Summarize: {text}", self.cur_retry==0)
+    def exec(self, prep_res):
+        return call_llm(f"Summarize: {prep_res}", self.cur_retry==0)
 ```
 
 - Enable logging:
diff --git a/.gitignore b/.gitignore
index 659af27f..7db945c5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,6 +15,7 @@ Thumbs.db
 *.swo
 *~
 
+
 # Node
 node_modules/
 npm-debug.log
@@ -49,6 +50,8 @@ wheels/
 *.egg
 venv/
 ENV/
+.venv/
+.ruff_cache/
 
 # Logs and databases
 *.log
diff --git a/README.md b/README.md
index efaa8f81..47c0771a 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,8 @@ English | [中文](https://github.com/The-Pocket/PocketFlow/blob/main/cookbook/p
     <img src="https://img.shields.io/discord/1346833819172601907?logo=discord&style=flat">
 </a>
 
+**! Edit: This fork has type hints and docstrings added for a total of 811 lines !**
+
 Pocket Flow is a [100-line](https://github.com/The-Pocket/PocketFlow/blob/main/pocketflow/__init__.py) minimalist LLM framework
 
 - **Lightweight**: Just 100 lines. Zero bloat, zero dependencies, zero vendor lock-in.
diff --git a/cookbook/pocketflow-a2a/nodes.py b/cookbook/pocketflow-a2a/nodes.py
index d777bdeb..8b9e5cb4 100644
--- a/cookbook/pocketflow-a2a/nodes.py
+++ b/cookbook/pocketflow-a2a/nodes.py
@@ -1,6 +1,8 @@
+import yaml
+
 from pocketflow import Node
 from utils import call_llm, search_web
-import yaml
+
 
 class DecideAction(Node):
     def prep(self, shared):
@@ -12,11 +14,11 @@ def prep(self, shared):
         # Return both for the exec step
         return question, context
         
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Call the LLM to decide whether to search or answer."""
-        question, context = inputs
+        question, context = prep_res
         
-        print(f"🤔 Agent deciding what to do next...")
+        print("🤔 Agent deciding what to do next...")
         
         # Create a prompt to help the LLM decide what to do next with proper yaml formatting
         prompt = f"""
@@ -71,7 +73,7 @@ def post(self, shared, prep_res, exec_res):
             print(f"🔍 Agent decided to search for: {exec_res['search_query']}")
         else:
             shared["context"] = exec_res["answer"] #save the context if LLM gives the answer without searching.
-            print(f"💡 Agent decided to answer the question")
+            print("💡 Agent decided to answer the question")
         
         # Return the action to determine the next node in the flow
         return exec_res["action"]
@@ -81,11 +83,11 @@ def prep(self, shared):
         """Get the search query from the shared store."""
         return shared["search_query"]
         
-    def exec(self, search_query):
+    def exec(self, prep_res):
         """Search the web for the given query."""
         # Call the search utility function
-        print(f"🌐 Searching the web for: {search_query}")
-        results = search_web(search_query)
+        print(f"🌐 Searching the web for: {prep_res}")
+        results = search_web(prep_res)
         return results
     
     def post(self, shared, prep_res, exec_res):
@@ -94,7 +96,7 @@ def post(self, shared, prep_res, exec_res):
         previous = shared.get("context", "")
         shared["context"] = previous + "\n\nSEARCH: " + shared["search_query"] + "\nRESULTS: " + exec_res
         
-        print(f"📚 Found information, analyzing results...")
+        print("📚 Found information, analyzing results...")
         
         # Always go back to the decision node after searching
         return "decide"
@@ -104,11 +106,11 @@ def prep(self, shared):
         """Get the question and context for answering."""
         return shared["question"], shared.get("context", "")
         
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Call the LLM to generate a final answer."""
-        question, context = inputs
+        question, context = prep_res
         
-        print(f"✍️ Crafting final answer...")
+        print("✍️ Crafting final answer...")
         
         # Create a prompt for the LLM to answer the question
         prompt = f"""
@@ -129,7 +131,7 @@ def post(self, shared, prep_res, exec_res):
         # Save the answer in the shared store
         shared["answer"] = exec_res
         
-        print(f"✅ Answer generated successfully")
+        print("✅ Answer generated successfully")
         
         # We're done - no need to continue the flow
         return "done" 
diff --git a/cookbook/pocketflow-agent/demo.ipynb b/cookbook/pocketflow-agent/demo.ipynb
index fa155296..5dbc594f 100644
--- a/cookbook/pocketflow-agent/demo.ipynb
+++ b/cookbook/pocketflow-agent/demo.ipynb
@@ -4,10 +4,10 @@
       "cell_type": "code",
       "execution_count": 1,
       "metadata": {
+        "id": "8MeqVASIxKBH",
         "vscode": {
           "languageId": "plaintext"
-        },
-        "id": "8MeqVASIxKBH"
+        }
       },
       "outputs": [],
       "source": [
@@ -21,19 +21,19 @@
       "cell_type": "code",
       "execution_count": 2,
       "metadata": {
-        "vscode": {
-          "languageId": "plaintext"
-        },
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "wUp_sNU1xKBI",
-        "outputId": "a647f919-b253-48c8-c132-5eef582e29c7"
+        "outputId": "a647f919-b253-48c8-c132-5eef582e29c7",
+        "vscode": {
+          "languageId": "plaintext"
+        }
       },
       "outputs": [
         {
-          "output_type": "stream",
           "name": "stdout",
+          "output_type": "stream",
           "text": [
             "## Testing call_llm\n",
             "## Prompt: In a few words, what is the meaning of life?\n",
@@ -97,12 +97,12 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": null,
       "metadata": {
+        "id": "T0ETd4C2xKBI",
         "vscode": {
           "languageId": "plaintext"
-        },
-        "id": "T0ETd4C2xKBI"
+        }
       },
       "outputs": [],
       "source": [
@@ -120,9 +120,9 @@
         "        # Return both for the exec step\n",
         "        return question, context\n",
         "\n",
-        "    def exec(self, inputs):\n",
+        "    def exec(self, prep_res):\n",
         "        \"\"\"Call the LLM to decide whether to search or answer.\"\"\"\n",
-        "        question, context = inputs\n",
+        "        question, context = prep_res\n",
         "\n",
         "        print(f\"🤔 Agent deciding what to do next...\")\n",
         "\n",
@@ -189,11 +189,11 @@
         "        \"\"\"Get the search query from the shared store.\"\"\"\n",
         "        return shared[\"search_query\"]\n",
         "\n",
-        "    def exec(self, search_query):\n",
+        "    def exec(self, prep_res):\n",
         "        \"\"\"Search the web for the given query.\"\"\"\n",
         "        # Call the search utility function\n",
-        "        print(f\"🌐 Searching the web for: {search_query}\")\n",
-        "        results = search_web(search_query)\n",
+        "        print(f\"🌐 Searching the web for: {prep_res}\")\n",
+        "        results = search_web(prep_res)\n",
         "        return results\n",
         "\n",
         "    def post(self, shared, prep_res, exec_res):\n",
@@ -212,9 +212,9 @@
         "        \"\"\"Get the question and context for answering.\"\"\"\n",
         "        return shared[\"question\"], shared.get(\"context\", \"\")\n",
         "\n",
-        "    def exec(self, inputs):\n",
+        "    def exec(self, prep_res):\n",
         "        \"\"\"Call the LLM to generate a final answer.\"\"\"\n",
-        "        question, context = inputs\n",
+        "        question, context = prep_res\n",
         "\n",
         "        print(f\"✍️ Crafting final answer...\")\n",
         "\n",
@@ -247,10 +247,10 @@
       "cell_type": "code",
       "execution_count": 4,
       "metadata": {
+        "id": "0B4jCAmXxKBI",
         "vscode": {
           "languageId": "plaintext"
-        },
-        "id": "0B4jCAmXxKBI"
+        }
       },
       "outputs": [],
       "source": [
@@ -293,19 +293,19 @@
       "cell_type": "code",
       "execution_count": 5,
       "metadata": {
-        "vscode": {
-          "languageId": "plaintext"
-        },
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "bIwsNEDCxKBI",
-        "outputId": "e6c02020-6fae-4377-8f0a-01d2580dd659"
+        "outputId": "e6c02020-6fae-4377-8f0a-01d2580dd659",
+        "vscode": {
+          "languageId": "plaintext"
+        }
       },
       "outputs": [
         {
-          "output_type": "stream",
           "name": "stdout",
+          "output_type": "stream",
           "text": [
             "🤔 Processing question: Who won the Nobel Prize in Physics 2024?\n",
             "🤔 Agent deciding what to do next...\n",
@@ -353,15 +353,15 @@
     }
   ],
   "metadata": {
-    "language_info": {
-      "name": "python"
-    },
     "colab": {
       "provenance": []
     },
     "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
     }
   },
   "nbformat": 4,
diff --git a/cookbook/pocketflow-agent/nodes.py b/cookbook/pocketflow-agent/nodes.py
index d2fd2024..ceff2bbb 100644
--- a/cookbook/pocketflow-agent/nodes.py
+++ b/cookbook/pocketflow-agent/nodes.py
@@ -1,6 +1,8 @@
+import yaml
+
 from pocketflow import Node
 from utils import call_llm, search_web_duckduckgo
-import yaml
+
 
 class DecideAction(Node):
     def prep(self, shared):
@@ -12,11 +14,11 @@ def prep(self, shared):
         # Return both for the exec step
         return question, context
         
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Call the LLM to decide whether to search or answer."""
-        question, context = inputs
+        question, context = prep_res
         
-        print(f"🤔 Agent deciding what to do next...")
+        print("🤔 Agent deciding what to do next...")
         
         # Create a prompt to help the LLM decide what to do next with proper yaml formatting
         prompt = f"""
@@ -71,7 +73,7 @@ def post(self, shared, prep_res, exec_res):
             print(f"🔍 Agent decided to search for: {exec_res['search_query']}")
         else:
             shared["context"] = exec_res["answer"] #save the context if LLM gives the answer without searching.
-            print(f"💡 Agent decided to answer the question")
+            print("💡 Agent decided to answer the question")
         
         # Return the action to determine the next node in the flow
         return exec_res["action"]
@@ -81,11 +83,11 @@ def prep(self, shared):
         """Get the search query from the shared store."""
         return shared["search_query"]
         
-    def exec(self, search_query):
+    def exec(self, prep_res):
         """Search the web for the given query."""
         # Call the search utility function
-        print(f"🌐 Searching the web for: {search_query}")
-        results = search_web_duckduckgo(search_query)
+        print(f"🌐 Searching the web for: {prep_res}")
+        results = search_web_duckduckgo(prep_res)
         return results
     
     def post(self, shared, prep_res, exec_res):
@@ -94,7 +96,7 @@ def post(self, shared, prep_res, exec_res):
         previous = shared.get("context", "")
         shared["context"] = previous + "\n\nSEARCH: " + shared["search_query"] + "\nRESULTS: " + exec_res
         
-        print(f"📚 Found information, analyzing results...")
+        print("📚 Found information, analyzing results...")
         
         # Always go back to the decision node after searching
         return "decide"
@@ -104,11 +106,11 @@ def prep(self, shared):
         """Get the question and context for answering."""
         return shared["question"], shared.get("context", "")
         
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Call the LLM to generate a final answer."""
-        question, context = inputs
+        question, context = prep_res
         
-        print(f"✍️ Crafting final answer...")
+        print("✍️ Crafting final answer...")
         
         # Create a prompt for the LLM to answer the question
         prompt = f"""
@@ -129,7 +131,7 @@ def post(self, shared, prep_res, exec_res):
         # Save the answer in the shared store
         shared["answer"] = exec_res
         
-        print(f"✅ Answer generated successfully")
+        print("✅ Answer generated successfully")
         
         # We're done - no need to continue the flow
         return "done" 
diff --git a/cookbook/pocketflow-async-basic/README.md b/cookbook/pocketflow-async-basic/README.md
index 5928f5b1..6697aeeb 100644
--- a/cookbook/pocketflow-async-basic/README.md
+++ b/cookbook/pocketflow-async-basic/README.md
@@ -22,18 +22,18 @@ When you run the example:
        ingredient = input("Enter ingredient: ")
        return ingredient
 
-   async def exec_async(self, ingredient):
+   async def exec_async(self, prep_res):
        # Async API call
-       recipes = await fetch_recipes(ingredient)
+       recipes = await fetch_recipes(prep_res)
        return recipes
    ```
 
 2. **SuggestRecipe (AsyncNode)**
    ```python
-   async def exec_async(self, recipes):
+   async def exec_async(self, prep_res):
        # Async LLM call
        suggestion = await call_llm_async(
-           f"Choose best recipe from: {recipes}"
+           f"Choose best recipe from: {prep_res}"
        )
        return suggestion
    ```
diff --git a/cookbook/pocketflow-async-basic/nodes.py b/cookbook/pocketflow-async-basic/nodes.py
index 281dd79e..c3abb876 100644
--- a/cookbook/pocketflow-async-basic/nodes.py
+++ b/cookbook/pocketflow-async-basic/nodes.py
@@ -1,5 +1,6 @@
 from pocketflow import AsyncNode
-from utils import fetch_recipes, call_llm_async, get_user_input
+from utils import call_llm_async, fetch_recipes, get_user_input
+
 
 class FetchRecipes(AsyncNode):
     """AsyncNode that fetches recipes."""
@@ -9,14 +10,14 @@ async def prep_async(self, shared):
         ingredient = await get_user_input("Enter ingredient: ")
         return ingredient
     
-    async def exec_async(self, ingredient):
+    async def exec_async(self, prep_res):
         """Fetch recipes asynchronously."""
-        recipes = await fetch_recipes(ingredient)
+        recipes = await fetch_recipes(prep_res)
         return recipes
     
-    async def post_async(self, shared, prep_res, recipes):
+    async def post_async(self, shared, prep_res, exec_res):
         """Store recipes and continue."""
-        shared["recipes"] = recipes
+        shared["recipes"] = exec_res
         shared["ingredient"] = prep_res
         return "suggest"
 
@@ -27,16 +28,16 @@ async def prep_async(self, shared):
         """Get recipes from shared store."""
         return shared["recipes"]
     
-    async def exec_async(self, recipes):
+    async def exec_async(self, prep_res):
         """Get suggestion from LLM."""
         suggestion = await call_llm_async(
-            f"Choose best recipe from: {', '.join(recipes)}"
+            f"Choose best recipe from: {', '.join(prep_res)}"
         )
         return suggestion
     
-    async def post_async(self, shared, prep_res, suggestion):
+    async def post_async(self, shared, prep_res, exec_res):
         """Store suggestion and continue."""
-        shared["suggestion"] = suggestion
+        shared["suggestion"] = exec_res
         return "approve"
 
 class GetApproval(AsyncNode):
@@ -46,14 +47,14 @@ async def prep_async(self, shared):
         """Get current suggestion."""
         return shared["suggestion"]
     
-    async def exec_async(self, suggestion):
+    async def exec_async(self, prep_res):
         """Ask for user approval."""
-        answer = await get_user_input(f"\nAccept this recipe? (y/n): ")
+        answer = await get_user_input("\nAccept this recipe? (y/n): ")
         return answer
     
-    async def post_async(self, shared, prep_res, answer):
+    async def post_async(self, shared, prep_res, exec_res):
         """Handle user's decision."""
-        if answer == "y":
+        if exec_res == "y":
             print("\nGreat choice! Here's your recipe...")
             print(f"Recipe: {shared['suggestion']}")
             print(f"Ingredient: {shared['ingredient']}")
diff --git a/cookbook/pocketflow-batch-flow/nodes.py b/cookbook/pocketflow-batch-flow/nodes.py
index fb14a947..8c44b12e 100644
--- a/cookbook/pocketflow-batch-flow/nodes.py
+++ b/cookbook/pocketflow-batch-flow/nodes.py
@@ -1,19 +1,22 @@
 """Node implementations for image processing."""
 
 import os
+
 from PIL import Image, ImageEnhance, ImageFilter
+
 from pocketflow import Node
 
+
 class LoadImage(Node):
     """Node that loads an image file."""
     
     def prep(self, shared):
         """Get image path from parameters."""
-        return os.path.join("images", self.params["input"])
+        return os.path.join("images", str(self.params["input"]))
     
-    def exec(self, image_path):
+    def exec(self, prep_res):
         """Load the image using PIL."""
-        return Image.open(image_path)
+        return Image.open(prep_res)
     
     def post(self, shared, prep_res, exec_res):
         """Store the image in shared store."""
@@ -27,9 +30,9 @@ def prep(self, shared):
         """Get image and filter type."""
         return shared["image"], self.params["filter"]
     
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Apply the specified filter."""
-        image, filter_type = inputs
+        image, filter_type = prep_res
         
         if filter_type == "grayscale":
             return image.convert("L")
@@ -58,15 +61,15 @@ def prep(self, shared):
         os.makedirs("output", exist_ok=True)
         
         # Generate output filename
-        input_name = os.path.splitext(self.params["input"])[0]
+        input_name = os.path.splitext(str(self.params["input"]))[0]
         filter_name = self.params["filter"]
         output_path = os.path.join("output", f"{input_name}_{filter_name}.jpg")
         
         return shared["filtered_image"], output_path
     
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Save the image to file."""
-        image, output_path = inputs
+        image, output_path = prep_res
         image.save(output_path, "JPEG")
         return output_path
     
diff --git a/cookbook/pocketflow-batch-node/nodes.py b/cookbook/pocketflow-batch-node/nodes.py
index 7d5d4951..223b1ddd 100644
--- a/cookbook/pocketflow-batch-node/nodes.py
+++ b/cookbook/pocketflow-batch-node/nodes.py
@@ -1,6 +1,8 @@
 import pandas as pd
+
 from pocketflow import BatchNode
 
+
 class CSVProcessor(BatchNode):
     """BatchNode that processes a large CSV file in chunks."""
     
@@ -21,7 +23,7 @@ def prep(self, shared):
         )
         return chunks
     
-    def exec(self, chunk):
+    def exec(self, prep_res):
         """Process a single chunk of the CSV.
         
         Args:
@@ -31,25 +33,25 @@ def exec(self, chunk):
             dict: Statistics for this chunk
         """
         return {
-            "total_sales": chunk["amount"].sum(),
-            "num_transactions": len(chunk),
-            "total_amount": chunk["amount"].sum()
+            "total_sales": prep_res["amount"].sum(),
+            "num_transactions": len(prep_res),
+            "total_amount": prep_res["amount"].sum()
         }
     
-    def post(self, shared, prep_res, exec_res_list):
+    def post(self, shared, prep_res, exec_res):
         """Combine results from all chunks.
         
         Args:
             prep_res: Original chunks iterator
-            exec_res_list: List of results from each chunk
+            exec_res: List of results from each chunk
             
         Returns:
             str: Action to take next
         """
         # Combine statistics from all chunks
-        total_sales = sum(res["total_sales"] for res in exec_res_list)
-        total_transactions = sum(res["num_transactions"] for res in exec_res_list)
-        total_amount = sum(res["total_amount"] for res in exec_res_list)
+        total_sales = sum(res["total_sales"] for res in exec_res)
+        total_transactions = sum(res["num_transactions"] for res in exec_res)
+        total_amount = sum(res["total_amount"] for res in exec_res)
         
         # Calculate final statistics
         shared["statistics"] = {
diff --git a/cookbook/pocketflow-batch/main.py b/cookbook/pocketflow-batch/main.py
index cc0b79b0..0f56c684 100644
--- a/cookbook/pocketflow-batch/main.py
+++ b/cookbook/pocketflow-batch/main.py
@@ -1,8 +1,10 @@
 import os
 import time
+
 from pocketflow import BatchNode, Flow
 from utils import call_llm
 
+
 class TranslateTextNode(BatchNode):
     def prep(self, shared):
         text = shared.get("text", "(No text provided)")
@@ -12,8 +14,8 @@ def prep(self, shared):
         # Create batches for each language translation
         return [(text, lang) for lang in languages]
 
-    def exec(self, data_tuple):
-        text, language = data_tuple
+    def exec(self, prep_res):
+        text, language = prep_res
         
         prompt = f"""
 Please translate the following markdown file into {language}. 
@@ -29,13 +31,13 @@ def exec(self, data_tuple):
         print(f"Translated {language} text")
         return {"language": language, "translation": result}
 
-    def post(self, shared, prep_res, exec_res_list):
+    def post(self, shared, prep_res, exec_res):
         # Create output directory if it doesn't exist
         output_dir = shared.get("output_dir", "translations")
         os.makedirs(output_dir, exist_ok=True)
         
         # Write each translation to a file
-        for result in exec_res_list:
+        for result in exec_res:
             language, translation = result["language"], result["translation"]
             
             # Write to file
diff --git a/cookbook/pocketflow-chat-guardrail/main.py b/cookbook/pocketflow-chat-guardrail/main.py
index 54bf6cb9..896fe7eb 100644
--- a/cookbook/pocketflow-chat-guardrail/main.py
+++ b/cookbook/pocketflow-chat-guardrail/main.py
@@ -1,6 +1,7 @@
-from pocketflow import Node, Flow
+from pocketflow import Flow, Node
 from utils import call_llm
 
+
 class UserInputNode(Node):
     def prep(self, shared):
         # Initialize messages if this is the first run
@@ -10,7 +11,7 @@ def prep(self, shared):
         
         return None
 
-    def exec(self, _):
+    def exec(self, prep_res):
         # Get user input
         user_input = input("\nYou: ")
         return user_input
@@ -35,19 +36,19 @@ def prep(self, shared):
         user_input = shared.get("user_input", "")
         return user_input
     
-    def exec(self, user_input):
+    def exec(self, prep_res):
         # Basic validation checks
-        if not user_input or user_input.strip() == "":
+        if not prep_res or prep_res.strip() == "":
             return False, "Your query is empty. Please provide a travel-related question."
         
-        if len(user_input.strip()) < 3:
+        if len(prep_res.strip()) < 3:
             return False, "Your query is too short. Please provide more details about your travel question."
         
         # LLM-based validation for travel topics
         prompt = f"""
 Evaluate if the following user query is related to travel advice, destinations, planning, or other travel topics.
 The chat should ONLY answer travel-related questions and reject any off-topic, harmful, or inappropriate queries.
-User query: {user_input}
+User query: {prep_res}
 Return your evaluation in YAML format:
 ```yaml
 valid: true/false
@@ -96,9 +97,9 @@ def prep(self, shared):
         # Return all messages for the LLM
         return shared["messages"]
 
-    def exec(self, messages):
+    def exec(self, prep_res):
         # Call LLM with the entire conversation history
-        response = call_llm(messages)
+        response = call_llm(prep_res)
         return response
 
     def post(self, shared, prep_res, exec_res):
diff --git a/cookbook/pocketflow-chat-memory/nodes.py b/cookbook/pocketflow-chat-memory/nodes.py
index 02237db2..ef7f56a1 100644
--- a/cookbook/pocketflow-chat-memory/nodes.py
+++ b/cookbook/pocketflow-chat-memory/nodes.py
@@ -1,7 +1,8 @@
 from pocketflow import Node
-from utils.vector_index import create_index, add_vector, search_vectors
 from utils.call_llm import call_llm
 from utils.get_embedding import get_embedding
+from utils.vector_index import add_vector, create_index, search_vectors
+
 
 class GetUserQuestionNode(Node):
     def prep(self, shared):
@@ -12,7 +13,7 @@ def prep(self, shared):
         
         return None
     
-    def exec(self, _):
+    def exec(self, prep_res):
         """Get user input interactively"""
         # Get interactive input from user
         user_input = input("\nYou: ")
@@ -62,13 +63,13 @@ def prep(self, shared):
         
         return context
     
-    def exec(self, messages):
+    def exec(self, prep_res):
         """Generate a response using the LLM"""
-        if messages is None:
+        if prep_res is None:
             return None
         
         # Call LLM with the context
-        response = call_llm(messages)
+        response = call_llm(prep_res)
         return response
     
     def post(self, shared, prep_res, exec_res):
@@ -103,21 +104,21 @@ def prep(self, shared):
         
         return oldest_pair
     
-    def exec(self, conversation):
+    def exec(self, prep_res):
         """Embed a conversation"""
-        if not conversation:
+        if not prep_res:
             return None
             
         # Combine user and assistant messages into a single text for embedding
-        user_msg = next((msg for msg in conversation if msg["role"] == "user"), {"content": ""})
-        assistant_msg = next((msg for msg in conversation if msg["role"] == "assistant"), {"content": ""})
+        user_msg = next((msg for msg in prep_res if msg["role"] == "user"), {"content": ""})
+        assistant_msg = next((msg for msg in prep_res if msg["role"] == "assistant"), {"content": ""})
         combined = f"User: {user_msg['content']} Assistant: {assistant_msg['content']}"
         
         # Generate embedding
         embedding = get_embedding(combined)
         
         return {
-            "conversation": conversation,
+            "conversation": prep_res,
             "embedding": embedding
         }
     
@@ -164,14 +165,14 @@ def prep(self, shared):
             "vector_items": shared["vector_items"]
         }
     
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Find the most relevant past conversation"""
-        if not inputs:
+        if not prep_res:
             return None
             
-        query = inputs["query"]
-        vector_index = inputs["vector_index"]
-        vector_items = inputs["vector_items"]
+        query = prep_res["query"]
+        vector_index = prep_res["vector_index"]
+        vector_items = prep_res["vector_items"]
         
         print(f"🔍 Finding relevant conversation for: {query[:30]}...")
         
diff --git a/cookbook/pocketflow-chat/main.py b/cookbook/pocketflow-chat/main.py
index 126c1064..ab14c7fb 100644
--- a/cookbook/pocketflow-chat/main.py
+++ b/cookbook/pocketflow-chat/main.py
@@ -1,6 +1,7 @@
-from pocketflow import Node, Flow
+from pocketflow import Flow, Node
 from utils import call_llm
 
+
 class ChatNode(Node):
     def prep(self, shared):
         # Initialize messages if this is the first run
@@ -21,12 +22,12 @@ def prep(self, shared):
         # Return all messages for the LLM
         return shared["messages"]
 
-    def exec(self, messages):
-        if messages is None:
+    def exec(self, prep_res):
+        if prep_res is None:
             return None
         
         # Call LLM with the entire conversation history
-        response = call_llm(messages)
+        response = call_llm(prep_res)
         return response
 
     def post(self, shared, prep_res, exec_res):
diff --git a/cookbook/pocketflow-cli-hitl/nodes.py b/cookbook/pocketflow-cli-hitl/nodes.py
index ea2bf818..1221bdef 100644
--- a/cookbook/pocketflow-cli-hitl/nodes.py
+++ b/cookbook/pocketflow-cli-hitl/nodes.py
@@ -1,12 +1,13 @@
 from pocketflow import Node
 from utils.call_llm import call_llm
 
+
 class GetTopicNode(Node):
     """Prompts the user to enter the topic for the joke."""
-    def exec(self, _shared):
+    def exec(self, prep_res):
         return input("What topic would you like a joke about? ")
 
-    def post(self, shared, _prep_res, exec_res):
+    def post(self, shared, prep_res, exec_res):
         shared["topic"] = exec_res
 
 class GenerateJokeNode(Node):
@@ -24,20 +25,20 @@ def prep(self, shared):
     def exec(self, prep_res):
         return call_llm(prep_res)
 
-    def post(self, shared, _prep_res, exec_res):
+    def post(self, shared, prep_res, exec_res):
         shared["current_joke"] = exec_res
         print(f"\nJoke: {exec_res}")
 
 class GetFeedbackNode(Node):
     """Presents the joke to the user and asks for approval."""
-    def exec(self, _prep_res):
+    def exec(self, prep_res):
         while True:
             feedback = input("Did you like this joke? (yes/no): ").strip().lower()
             if feedback in ["yes", "y", "no", "n"]:
                 return feedback
             print("Invalid input. Please type 'yes' or 'no'.")
 
-    def post(self, shared, _prep_res, exec_res):
+    def post(self, shared, prep_res, exec_res):
         if exec_res in ["yes", "y"]:
             shared["user_feedback"] = "approve"
             print("Great! Glad you liked it.")
diff --git a/cookbook/pocketflow-code-generator/nodes.py b/cookbook/pocketflow-code-generator/nodes.py
index a29b6e96..0b3e3437 100644
--- a/cookbook/pocketflow-code-generator/nodes.py
+++ b/cookbook/pocketflow-code-generator/nodes.py
@@ -1,16 +1,18 @@
 import yaml
-from pocketflow import Node, BatchNode
+
+from pocketflow import BatchNode, Node
 from utils.call_llm import call_llm
 from utils.code_executor import execute_python
 
+
 class GenerateTestCases(Node):
     def prep(self, shared):
         return shared["problem"]
 
-    def exec(self, problem):
+    def exec(self, prep_res):
         prompt = f"""Generate 5-7 test cases for this coding problem:
 
-{problem}
+{prep_res}
 
 Output in this YAML format with reasoning:
 ```yaml
@@ -57,8 +59,8 @@ class ImplementFunction(Node):
     def prep(self, shared):
         return shared["problem"], shared["test_cases"]
 
-    def exec(self, inputs):
-        problem, test_cases = inputs
+    def exec(self, prep_res):
+        problem, test_cases = prep_res
         
         # Format test cases nicely for the prompt
         formatted_tests = ""
@@ -101,7 +103,7 @@ def post(self, shared, prep_res, exec_res):
         shared["function_code"] = exec_res
         
         # Print the implemented function
-        print(f"\n=== Implemented Function ===")
+        print("\n=== Implemented Function ===")
         print(exec_res)
 
 class RunTests(BatchNode):
@@ -111,8 +113,8 @@ def prep(self, shared):
         # Return list of tuples (function_code, test_case)
         return [(function_code, test_case) for test_case in test_cases]
 
-    def exec(self, test_data):
-        function_code, test_case = test_data
+    def exec(self, prep_res):
+        function_code, test_case = prep_res
         output, error = execute_python(function_code, test_case["input"])
         
         if error:
@@ -133,17 +135,17 @@ def exec(self, test_data):
             "error": None if passed else f"Expected {test_case['expected']}, got {output}"
         }
 
-    def post(self, shared, prep_res, exec_res_list):
-        shared["test_results"] = exec_res_list
-        all_passed = all(result["passed"] for result in exec_res_list)
+    def post(self, shared, prep_res, exec_res):
+        shared["test_results"] = exec_res
+        all_passed = all(result["passed"] for result in exec_res)
         shared["iteration_count"] = shared.get("iteration_count", 0) + 1
         
         # Print test results
-        passed_count = len([r for r in exec_res_list if r["passed"]])
-        total_count = len(exec_res_list)
+        passed_count = len([r for r in exec_res if r["passed"]])
+        total_count = len(exec_res)
         print(f"\n=== Test Results: {passed_count}/{total_count} Passed ===")
         
-        failed_tests = [r for r in exec_res_list if not r["passed"]]
+        failed_tests = [r for r in exec_res if not r["passed"]]
         if failed_tests:
             print("Failed tests:")
             for i, result in enumerate(failed_tests, 1):
@@ -172,17 +174,17 @@ def prep(self, shared):
             "failed_tests": failed_tests
         }
 
-    def exec(self, inputs):
+    def exec(self, prep_res):
         # Format current test cases nicely
         formatted_tests = ""
-        for i, test in enumerate(inputs['test_cases'], 1):
+        for i, test in enumerate(prep_res['test_cases'], 1):
             formatted_tests += f"{i}. {test['name']}\n"
             formatted_tests += f"   input: {test['input']}\n"
             formatted_tests += f"   expected: {test['expected']}\n\n"
         
         # Format failed tests nicely
         formatted_failures = ""
-        for i, result in enumerate(inputs['failed_tests'], 1):
+        for i, result in enumerate(prep_res['failed_tests'], 1):
             test_case = result['test_case']
             formatted_failures += f"{i}. {test_case['name']}:\n"
             if result['error']:
@@ -191,14 +193,14 @@ def exec(self, inputs):
                 formatted_failures += f"   output: {result['actual']}\n"
             formatted_failures += f"   expected: {result['expected']}\n\n"
 
-        prompt = f"""Problem: {inputs['problem']}
+        prompt = f"""Problem: {prep_res['problem']}
 
 Current test cases:
 {formatted_tests}
 
 Current function:
 ```python
-{inputs['function_code']}
+{prep_res['function_code']}
 ```
 
 Failed tests:
diff --git a/cookbook/pocketflow-communication/nodes.py b/cookbook/pocketflow-communication/nodes.py
index 6e64d380..c7052098 100644
--- a/cookbook/pocketflow-communication/nodes.py
+++ b/cookbook/pocketflow-communication/nodes.py
@@ -2,6 +2,7 @@
 
 from pocketflow import Node
 
+
 class EndNode(Node):
     """Node that handles flow termination."""
     pass
@@ -38,9 +39,9 @@ def prep(self, shared):
         """Get text from shared store."""
         return shared["text"]
     
-    def exec(self, text):
+    def exec(self, prep_res):
         """Count words in the text."""
-        return len(text.split())
+        return len(prep_res.split())
     
     def post(self, shared, prep_res, exec_res):
         """Update word count statistics."""
@@ -57,7 +58,7 @@ def prep(self, shared):
     def post(self, shared, prep_res, exec_res):
         """Display statistics and continue the flow."""
         stats = prep_res
-        print(f"\nStatistics:")
+        print("\nStatistics:")
         print(f"- Texts processed: {stats['total_texts']}")
         print(f"- Total words: {stats['total_words']}")
         print(f"- Average words per text: {stats['total_words'] / stats['total_texts']:.1f}\n")
diff --git a/cookbook/pocketflow-fastapi-background/nodes.py b/cookbook/pocketflow-fastapi-background/nodes.py
index a7f4cca8..491a4798 100644
--- a/cookbook/pocketflow-fastapi-background/nodes.py
+++ b/cookbook/pocketflow-fastapi-background/nodes.py
@@ -1,14 +1,16 @@
 import yaml
-from pocketflow import Node, BatchNode
+
+from pocketflow import BatchNode, Node
 from utils.call_llm import call_llm
 
+
 class GenerateOutline(Node):
     def prep(self, shared):
         return shared["topic"]
     
-    def exec(self, topic):
+    def exec(self, prep_res):
         prompt = f"""
-Create a simple outline for an article about {topic}.
+Create a simple outline for an article about {prep_res}.
 Include at most 3 main sections (no subsections).
 
 Output the sections in YAML format as shown below:
@@ -41,11 +43,11 @@ def prep(self, shared):
         self.sse_queue = shared["sse_queue"]
         return self.sections
     
-    def exec(self, section):
+    def exec(self, prep_res):
         prompt = f"""
 Write a short paragraph (MAXIMUM 100 WORDS) about this section:
 
-{section}
+{prep_res}
 
 Requirements:
 - Explain the idea in simple, easy-to-understand terms
@@ -56,7 +58,7 @@ def exec(self, section):
         content = call_llm(prompt)
         
         # Send progress update for this section
-        current_section_index = self.sections.index(section) if section in self.sections else 0
+        current_section_index = self.sections.index(prep_res) if prep_res in self.sections else 0
         total_sections = len(self.sections)
         
         # Progress from 33% (after outline) to 66% (before styling)
@@ -67,17 +69,17 @@ def exec(self, section):
             "step": "content", 
             "progress": section_progress, 
             "data": {
-                "section": section,
+                "section": prep_res,
                 "completed_sections": current_section_index + 1,
                 "total_sections": total_sections
             }
         }
         self.sse_queue.put_nowait(progress_msg)
         
-        return f"## {section}\n\n{content}\n"
+        return f"## {prep_res}\n\n{content}\n"
     
-    def post(self, shared, prep_res, exec_res_list):
-        draft = "\n".join(exec_res_list)
+    def post(self, shared, prep_res, exec_res):
+        draft = "\n".join(exec_res)
         shared["draft"] = draft
         return "default"
 
@@ -85,11 +87,11 @@ class ApplyStyle(Node):
     def prep(self, shared):
         return shared["draft"]
     
-    def exec(self, draft):
+    def exec(self, prep_res):
         prompt = f"""
 Rewrite the following draft in a conversational, engaging style:
 
-{draft}
+{prep_res}
 
 Make it:
 - Conversational and warm in tone
diff --git a/cookbook/pocketflow-fastapi-websocket/main.py b/cookbook/pocketflow-fastapi-websocket/main.py
index 8f8152c5..a4ff845e 100644
--- a/cookbook/pocketflow-fastapi-websocket/main.py
+++ b/cookbook/pocketflow-fastapi-websocket/main.py
@@ -1,7 +1,8 @@
 import json
+
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
 from flow import create_streaming_chat_flow
 
 app = FastAPI()
@@ -16,7 +17,7 @@ async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     
     # Initialize conversation history for this connection
-    shared_store = {
+    shared_storage = {
         "websocket": websocket,
         "conversation_history": []
     }
@@ -27,10 +28,10 @@ async def websocket_endpoint(websocket: WebSocket):
             message = json.loads(data)
             
             # Update only the current message, keep conversation history
-            shared_store["user_message"] = message.get("content", "")
+            shared_storage["user_message"] = message.get("content", "")
             
             flow = create_streaming_chat_flow()
-            await flow.run_async(shared_store)
+            await flow.run_async(shared_storage)
             
     except WebSocketDisconnect:
         pass
diff --git a/cookbook/pocketflow-flow/flow.py b/cookbook/pocketflow-flow/flow.py
index 8b8d03c2..663a2787 100644
--- a/cookbook/pocketflow-flow/flow.py
+++ b/cookbook/pocketflow-flow/flow.py
@@ -1,4 +1,5 @@
-from pocketflow import Node, Flow
+from pocketflow import Flow, Node
+
 
 class TextInput(Node):
     def prep(self, shared):
@@ -28,8 +29,8 @@ class TextTransform(Node):
     def prep(self, shared):
         return shared["text"], shared["choice"]
     
-    def exec(self, inputs):
-        text, choice = inputs
+    def exec(self, prep_res):
+        text, choice = prep_res
         
         if choice == "1":
             return text.upper()
diff --git a/cookbook/pocketflow-google-calendar/nodes.py b/cookbook/pocketflow-google-calendar/nodes.py
index 5ec6c8d9..5fed34de 100644
--- a/cookbook/pocketflow-google-calendar/nodes.py
+++ b/cookbook/pocketflow-google-calendar/nodes.py
@@ -1,6 +1,6 @@
 from pocketflow import Node
-from utils.google_calendar import create_event, list_events, list_calendar_lists
-from datetime import datetime, timedelta
+from utils.google_calendar import create_event, list_calendar_lists, list_events
+
 
 class CreateCalendarEventNode(Node):
     def prep(self, shared):
@@ -12,14 +12,14 @@ def prep(self, shared):
             'end_time': shared.get('event_end_time')
         }
     
-    def exec(self, event_data):
+    def exec(self, prep_res):
         """Creates a new calendar event."""
         try:
             event = create_event(
-                summary=event_data['summary'],
-                description=event_data['description'],
-                start_time=event_data['start_time'],
-                end_time=event_data['end_time']
+                summary=prep_res['summary'],
+                description=prep_res['description'],
+                start_time=prep_res['start_time'],
+                end_time=prep_res['end_time']
             )
             return {'success': True, 'event': event}
         except Exception as e:
@@ -41,10 +41,10 @@ def prep(self, shared):
             'days': shared.get('days_to_list', 7)
         }
     
-    def exec(self, params):
+    def exec(self, prep_res):
         """Lists calendar events."""
         try:
-            events = list_events(days=params['days'])
+            events = list_events(days=prep_res['days'])
             return {'success': True, 'events': events}
         except Exception as e:
             return {'success': False, 'error': str(e)}
@@ -63,7 +63,7 @@ def prep(self, shared):
         """No special preparation needed to list calendars."""
         return {}
 
-    def exec(self, params):
+    def exec(self, prep_res):
         """Lists all available calendars for the user."""
         try:
             calendars = list_calendar_lists()
diff --git a/cookbook/pocketflow-hello-world/flow.py b/cookbook/pocketflow-hello-world/flow.py
index 5276a016..c2ca35ae 100644
--- a/cookbook/pocketflow-hello-world/flow.py
+++ b/cookbook/pocketflow-hello-world/flow.py
@@ -1,6 +1,7 @@
-from pocketflow import Node, Flow
+from pocketflow import Flow, Node
 from utils.call_llm import call_llm
 
+
 # An example node and flow
 # Please replace this with your own node and flow
 class AnswerNode(Node):
@@ -8,8 +9,8 @@ def prep(self, shared):
         # Read question from shared
         return shared["question"]
     
-    def exec(self, question):
-        return call_llm(question)
+    def exec(self, prep_res):
+        return call_llm(prep_res)
     
     def post(self, shared, prep_res, exec_res):
         # Store the answer in shared
diff --git a/cookbook/pocketflow-majority-vote/main.py b/cookbook/pocketflow-majority-vote/main.py
index ad517031..12077cbe 100644
--- a/cookbook/pocketflow-majority-vote/main.py
+++ b/cookbook/pocketflow-majority-vote/main.py
@@ -10,10 +10,10 @@ def prep(self, shared):
         attempts_count = shared.get("num_tries", 3)
         return [question for _ in range(attempts_count)]
 
-    def exec(self, single_question: str):
+    def exec(self, prep_res: str):
         prompt = f"""
 You are a helpful assistant. Please answer the user's question below.
-Question: {single_question}
+Question: {prep_res}
 
 Return strictly using the following YAML structure:
 ```yaml
diff --git a/cookbook/pocketflow-map-reduce/nodes.py b/cookbook/pocketflow-map-reduce/nodes.py
index cf6d3bfd..089f9350 100644
--- a/cookbook/pocketflow-map-reduce/nodes.py
+++ b/cookbook/pocketflow-map-reduce/nodes.py
@@ -1,12 +1,15 @@
-from pocketflow import Node, BatchNode
-from utils import call_llm
-import yaml
 import os
 
+import yaml
+
+from pocketflow import BatchNode, Node
+from utils import call_llm
+
+
 class ReadResumesNode(Node):
     """Map phase: Read all resumes from the data directory into shared storage."""
     
-    def exec(self, _):
+    def exec(self, prep_res):
         resume_files = {}
         data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
         
@@ -29,9 +32,9 @@ class EvaluateResumesNode(BatchNode):
     def prep(self, shared):
         return list(shared["resumes"].items())
     
-    def exec(self, resume_item):
+    def exec(self, prep_res):
         """Evaluate a single resume."""
-        filename, content = resume_item
+        filename, content = prep_res
         
         prompt = f"""
 Evaluate the following resume and determine if the candidate qualifies for an advanced technical role.
@@ -60,8 +63,8 @@ def exec(self, resume_item):
         
         return (filename, result)
 
-    def post(self, shared, prep_res, exec_res_list):
-        shared["evaluations"] = {filename: result for filename, result in exec_res_list}
+    def post(self, shared, prep_res, exec_res):
+        shared["evaluations"] = {filename: result for filename, result in exec_res}
         return "default"
 
 
@@ -71,12 +74,12 @@ class ReduceResultsNode(Node):
     def prep(self, shared):
         return shared["evaluations"]
     
-    def exec(self, evaluations):
+    def exec(self, prep_res):
         qualified_count = 0
-        total_count = len(evaluations)
+        total_count = len(prep_res)
         qualified_candidates = []
         
-        for filename, evaluation in evaluations.items():
+        for filename, evaluation in prep_res.items():
             if evaluation.get("qualifies", False):
                 qualified_count += 1
                 qualified_candidates.append(evaluation.get("candidate_name", "Unknown"))
diff --git a/cookbook/pocketflow-mcp/main.py b/cookbook/pocketflow-mcp/main.py
index 03de8f28..043236f4 100644
--- a/cookbook/pocketflow-mcp/main.py
+++ b/cookbook/pocketflow-mcp/main.py
@@ -1,8 +1,11 @@
-from pocketflow import Node, Flow
-from utils import call_llm, get_tools, call_tool
-import yaml
 import sys
 
+import yaml
+
+from pocketflow import Flow, Node
+from utils import call_llm, call_tool, get_tools
+
+
 class GetToolsNode(Node):
     def prep(self, shared):
         """Initialize and get tools"""
@@ -10,9 +13,9 @@ def prep(self, shared):
         print("🔍 Getting available tools...")
         return "simple_server.py"
 
-    def exec(self, server_path):
+    def exec(self, prep_res):
         """Retrieve tools from the MCP server"""
-        tools = get_tools(server_path)
+        tools = get_tools(prep_res)
         return tools
 
     def post(self, shared, prep_res, exec_res):
@@ -73,10 +76,10 @@ def prep(self, shared):
 """
         return prompt
 
-    def exec(self, prompt):
+    def exec(self, prep_res):
         """Call LLM to process the question and decide which tool to use"""
         print("🤔 Analyzing question and deciding which tool to use...")
-        response = call_llm(prompt)
+        response = call_llm(prep_res)
         return response
 
     def post(self, shared, prep_res, exec_res):
@@ -103,9 +106,9 @@ def prep(self, shared):
         """Prepare tool execution parameters"""
         return shared["tool_name"], shared["parameters"]
 
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Execute the chosen tool"""
-        tool_name, parameters = inputs
+        tool_name, parameters = prep_res
         print(f"🔧 Executing tool '{tool_name}' with parameters: {parameters}")
         result = call_tool("simple_server.py", tool_name, parameters)
         return result
diff --git a/cookbook/pocketflow-multi-agent/main.py b/cookbook/pocketflow-multi-agent/main.py
index 40452529..f559eb4e 100644
--- a/cookbook/pocketflow-multi-agent/main.py
+++ b/cookbook/pocketflow-multi-agent/main.py
@@ -1,7 +1,9 @@
 import asyncio
-from pocketflow import AsyncNode, AsyncFlow
+
+from pocketflow import AsyncFlow, AsyncNode
 from utils import call_llm
 
+
 class AsyncHinter(AsyncNode):
     async def prep_async(self, shared):
         # Wait for message from guesser (or empty string at start)
@@ -10,10 +12,10 @@ async def prep_async(self, shared):
             return None
         return shared["target_word"], shared["forbidden_words"], shared.get("past_guesses", [])
 
-    async def exec_async(self, inputs):
-        if inputs is None:
+    async def exec_async(self, prep_res):
+        if prep_res is None:
             return None
-        target, forbidden, past_guesses = inputs
+        target, forbidden, past_guesses = prep_res
         prompt = f"Generate hint for '{target}'\nForbidden words: {forbidden}"
         if past_guesses:
             prompt += f"\nPrevious wrong guesses: {past_guesses}\nMake hint more specific."
@@ -36,8 +38,8 @@ async def prep_async(self, shared):
         hint = await shared["guesser_queue"].get()
         return hint, shared.get("past_guesses", [])
 
-    async def exec_async(self, inputs):
-        hint, past_guesses = inputs
+    async def exec_async(self, prep_res):
+        hint, past_guesses = prep_res
         prompt = f"Given hint: {hint}, past wrong guesses: {past_guesses}, make a new guess. Directly reply a single word:"
         guess = call_llm(prompt)
         print(f"Guesser: I guess it's - {guess}")
diff --git a/cookbook/pocketflow-nested-batch/nodes.py b/cookbook/pocketflow-nested-batch/nodes.py
index 89de7f17..1f1521a1 100644
--- a/cookbook/pocketflow-nested-batch/nodes.py
+++ b/cookbook/pocketflow-nested-batch/nodes.py
@@ -1,25 +1,27 @@
 import os
+
 from pocketflow import Node
 
+
 class LoadGrades(Node):
     """Node that loads grades from a student's file."""
     
     def prep(self, shared):
         """Get file path from parameters."""
-        class_name = self.params["class"]
-        student_file = self.params["student"]
+        class_name = str(self.params["class"])
+        student_file = str(self.params["student"])
         return os.path.join("school", class_name, student_file)
     
-    def exec(self, file_path):
+    def exec(self, prep_res):
         """Load and parse grades from file."""
-        with open(file_path, 'r') as f:
+        with open(prep_res, 'r') as f:
             # Each line is a grade
             grades = [float(line.strip()) for line in f]
         return grades
     
-    def post(self, shared, prep_res, grades):
+    def post(self, shared, prep_res, exec_res):
         """Store grades in shared store."""
-        shared["grades"] = grades
+        shared["grades"] = exec_res
         return "calculate"
 
 class CalculateAverage(Node):
@@ -29,11 +31,11 @@ def prep(self, shared):
         """Get grades from shared store."""
         return shared["grades"]
     
-    def exec(self, grades):
+    def exec(self, prep_res):
         """Calculate average."""
-        return sum(grades) / len(grades)
+        return sum(prep_res) / len(prep_res)
     
-    def post(self, shared, prep_res, average):
+    def post(self, shared, prep_res, exec_res):
         """Store and print result."""
         # Store in results dictionary
         if "results" not in shared:
@@ -45,8 +47,8 @@ def post(self, shared, prep_res, average):
         if class_name not in shared["results"]:
             shared["results"][class_name] = {}
             
-        shared["results"][class_name][student] = average
+        shared["results"][class_name][student] = exec_res
         
         # Print individual result
-        print(f"- {student}: Average = {average:.1f}")
+        print(f"- {student}: Average = {exec_res:.1f}")
         return "default" 
\ No newline at end of file
diff --git a/cookbook/pocketflow-parallel-batch-flow/nodes.py b/cookbook/pocketflow-parallel-batch-flow/nodes.py
index 7f05951d..205975a4 100644
--- a/cookbook/pocketflow-parallel-batch-flow/nodes.py
+++ b/cookbook/pocketflow-parallel-batch-flow/nodes.py
@@ -1,10 +1,13 @@
 """AsyncNode implementations for image processing."""
-import os
 import asyncio
-from PIL import Image, ImageFilter
+import os
+
 import numpy as np
+from PIL import Image, ImageFilter
+
 from pocketflow import AsyncNode
 
+
 class LoadImage(AsyncNode):
     """Node that loads an image from file."""
     async def prep_async(self, shared):
@@ -13,11 +16,11 @@ async def prep_async(self, shared):
         print(f"Loading image: {image_path}")
         return image_path
     
-    async def exec_async(self, image_path):
+    async def exec_async(self, prep_res):
         """Load image using PIL."""
         # Simulate I/O delay
         await asyncio.sleep(0.5)
-        return Image.open(image_path)
+        return Image.open(prep_res)
     
     async def post_async(self, shared, prep_res, exec_res):
         """Store image in shared store."""
@@ -33,9 +36,9 @@ async def prep_async(self, shared):
         print(f"Applying {filter_type} filter...")
         return image, filter_type
     
-    async def exec_async(self, inputs):
+    async def exec_async(self, prep_res):
         """Apply the specified filter."""
-        image, filter_type = inputs
+        image, filter_type = prep_res
         
         # Simulate processing delay
         await asyncio.sleep(0.5)
@@ -68,7 +71,7 @@ class SaveImage(AsyncNode):
     async def prep_async(self, shared):
         """Prepare output path."""
         image = shared["filtered_image"]
-        base_name = os.path.splitext(os.path.basename(self.params["image_path"]))[0]
+        base_name = os.path.splitext(os.path.basename(str(self.params["image_path"])))[0]
         filter_type = self.params["filter"]
         output_path = f"output/{base_name}_{filter_type}.jpg"
         
@@ -77,9 +80,9 @@ async def prep_async(self, shared):
         
         return image, output_path
     
-    async def exec_async(self, inputs):
+    async def exec_async(self, prep_res):
         """Save the image."""
-        image, output_path = inputs
+        image, output_path = prep_res
         
         # Simulate I/O delay
         await asyncio.sleep(0.5)
diff --git a/cookbook/pocketflow-parallel-batch/main.py b/cookbook/pocketflow-parallel-batch/main.py
index 5f944613..99a1647d 100644
--- a/cookbook/pocketflow-parallel-batch/main.py
+++ b/cookbook/pocketflow-parallel-batch/main.py
@@ -1,6 +1,7 @@
 import asyncio
-import time
 import os
+import time
+
 from pocketflow import AsyncFlow, AsyncParallelBatchNode
 from utils import call_llm
 
@@ -14,9 +15,9 @@ async def prep_async(self, shared):
         languages = shared.get("languages", [])
         return [(text, lang) for lang in languages]
 
-    async def exec_async(self, data_tuple):
+    async def exec_async(self, prep_res):
         """Calls the async LLM utility for each target language."""
-        text, language = data_tuple
+        text, language = prep_res
         
         prompt = f"""
 Please translate the following markdown file into {language}. 
diff --git a/cookbook/pocketflow-rag/nodes.py b/cookbook/pocketflow-rag/nodes.py
index e0b3cdda..92befba5 100644
--- a/cookbook/pocketflow-rag/nodes.py
+++ b/cookbook/pocketflow-rag/nodes.py
@@ -1,7 +1,9 @@
-from pocketflow import Node, Flow, BatchNode
-import numpy as np
 import faiss
-from utils import call_llm, get_embedding, fixed_size_chunk
+import numpy as np
+
+from pocketflow import BatchNode, Node
+from utils import call_llm, fixed_size_chunk, get_embedding
+
 
 # Nodes for the offline flow
 class ChunkDocumentsNode(BatchNode):
@@ -9,15 +11,15 @@ def prep(self, shared):
         """Read texts from shared store"""
         return shared["texts"]
     
-    def exec(self, text):
+    def exec(self, prep_res):
         """Chunk a single text into smaller pieces"""
-        return fixed_size_chunk(text)
+        return fixed_size_chunk(prep_res)
     
-    def post(self, shared, prep_res, exec_res_list):
+    def post(self, shared, prep_res, exec_res):
         """Store chunked texts in the shared store"""
         # Flatten the list of lists into a single list of chunks
         all_chunks = []
-        for chunks in exec_res_list:
+        for chunks in exec_res:
             all_chunks.extend(chunks)
         
         # Replace the original texts with the flat list of chunks
@@ -31,13 +33,13 @@ def prep(self, shared):
         """Read texts from shared store and return as an iterable"""
         return shared["texts"]
     
-    def exec(self, text):
+    def exec(self, prep_res):
         """Embed a single text"""
-        return get_embedding(text)
+        return get_embedding(prep_res)
     
-    def post(self, shared, prep_res, exec_res_list):
+    def post(self, shared, prep_res, exec_res):
         """Store embeddings in the shared store"""
-        embeddings = np.array(exec_res_list, dtype=np.float32)
+        embeddings = np.array(exec_res, dtype=np.float32)
         shared["embeddings"] = embeddings
         print(f"✅ Created {len(embeddings)} document embeddings")
         return "default"
@@ -47,16 +49,16 @@ def prep(self, shared):
         """Get embeddings from shared store"""
         return shared["embeddings"]
     
-    def exec(self, embeddings):
+    def exec(self, prep_res):
         """Create FAISS index and add embeddings"""
         print("🔍 Creating search index...")
-        dimension = embeddings.shape[1]
+        dimension = prep_res.shape[1]
         
         # Create a flat L2 index
         index = faiss.IndexFlatL2(dimension)
         
         # Add the embeddings to the index
-        index.add(embeddings)
+        index.add(prep_res)
         
         return index
     
@@ -72,10 +74,10 @@ def prep(self, shared):
         """Get query from shared store"""
         return shared["query"]
     
-    def exec(self, query):
+    def exec(self, prep_res):
         """Embed the query"""
-        print(f"🔍 Embedding query: {query}")
-        query_embedding = get_embedding(query)
+        print(f"🔍 Embedding query: {prep_res}")
+        query_embedding = get_embedding(prep_res)
         return np.array([query_embedding], dtype=np.float32)
     
     def post(self, shared, prep_res, exec_res):
@@ -88,10 +90,10 @@ def prep(self, shared):
         """Get query embedding, index, and texts from shared store"""
         return shared["query_embedding"], shared["index"], shared["texts"]
     
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Search the index for similar documents"""
         print("🔎 Searching for relevant documents...")
-        query_embedding, index, texts = inputs
+        query_embedding, index, texts = prep_res
         
         # Search for the most similar document
         distances, indices = index.search(query_embedding, k=1)
@@ -121,9 +123,9 @@ def prep(self, shared):
         """Get query, retrieved document, and any other context needed"""
         return shared["query"], shared["retrieved_document"]
     
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Generate an answer using the LLM"""
-        query, retrieved_doc = inputs
+        query, retrieved_doc = prep_res
         
         prompt = f"""
 Briefly answer the following question based on the context provided:
diff --git a/cookbook/pocketflow-streamlit-fsm/nodes.py b/cookbook/pocketflow-streamlit-fsm/nodes.py
index e5306967..d2254414 100644
--- a/cookbook/pocketflow-streamlit-fsm/nodes.py
+++ b/cookbook/pocketflow-streamlit-fsm/nodes.py
@@ -1,14 +1,15 @@
 from pocketflow import Node
 from utils.generate_image import generate_image
 
+
 class GenerateImageNode(Node):
     """Generates image from text prompt using OpenAI API."""
     
     def prep(self, shared):
         return shared.get("task_input", "")
 
-    def exec(self, prompt):
-        return generate_image(prompt)
+    def exec(self, prep_res):
+        return generate_image(prep_res)
 
     def post(self, shared, prep_res, exec_res):
         shared["input_used_by_process"] = prep_res
diff --git a/cookbook/pocketflow-supervisor/nodes.py b/cookbook/pocketflow-supervisor/nodes.py
index 99fbe870..a358148d 100644
--- a/cookbook/pocketflow-supervisor/nodes.py
+++ b/cookbook/pocketflow-supervisor/nodes.py
@@ -1,7 +1,10 @@
+import random
+
+import yaml
+
 from pocketflow import Node
 from utils import call_llm, search_web
-import yaml
-import random
+
 
 class DecideAction(Node):
     def prep(self, shared):
@@ -13,11 +16,11 @@ def prep(self, shared):
         # Return both for the exec step
         return question, context
         
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Call the LLM to decide whether to search or answer."""
-        question, context = inputs
+        question, context = prep_res
         
-        print(f"🤔 Agent deciding what to do next...")
+        print("🤔 Agent deciding what to do next...")
         
         # Create a prompt to help the LLM decide what to do next
         prompt = f"""
@@ -65,7 +68,7 @@ def post(self, shared, prep_res, exec_res):
             shared["search_query"] = exec_res["search_query"]
             print(f"🔍 Agent decided to search for: {exec_res['search_query']}")
         else:
-            print(f"💡 Agent decided to answer the question")
+            print("💡 Agent decided to answer the question")
         
         # Return the action to determine the next node in the flow
         return exec_res["action"]
@@ -75,11 +78,11 @@ def prep(self, shared):
         """Get the search query from the shared store."""
         return shared["search_query"]
         
-    def exec(self, search_query):
+    def exec(self, prep_res):
         """Search the web for the given query."""
         # Call the search utility function
-        print(f"🌐 Searching the web for: {search_query}")
-        results = search_web(search_query)
+        print(f"🌐 Searching the web for: {prep_res}")
+        results = search_web(prep_res)
         return results
     
     def post(self, shared, prep_res, exec_res):
@@ -88,7 +91,7 @@ def post(self, shared, prep_res, exec_res):
         previous = shared.get("context", "")
         shared["context"] = previous + "\n\nSEARCH: " + shared["search_query"] + "\nRESULTS: " + exec_res
         
-        print(f"📚 Found information, analyzing results...")
+        print("📚 Found information, analyzing results...")
         
         # Always go back to the decision node after searching
         return "decide"
@@ -98,16 +101,16 @@ def prep(self, shared):
         """Get the question and context for answering."""
         return shared["question"], shared.get("context", "")
         
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Call the LLM to generate a final answer with 50% chance of returning a dummy answer."""
-        question, context = inputs
+        question, context = prep_res
         
         # 50% chance to return a dummy answer
         if random.random() < 0.5:
-            print(f"🤪 Generating unreliable dummy answer...")
+            print("🤪 Generating unreliable dummy answer...")
             return "Sorry, I'm on a coffee break right now. All information I provide is completely made up anyway. The answer to your question is 42, or maybe purple unicorns. Who knows? Certainly not me!"
         
-        print(f"✍️ Crafting final answer...")
+        print("✍️ Crafting final answer...")
         
         # Create a prompt for the LLM to answer the question
         prompt = f"""
@@ -128,16 +131,16 @@ def post(self, shared, prep_res, exec_res):
         # Save the answer in the shared store
         shared["answer"] = exec_res
         
-        print(f"✅ Answer generated successfully")
+        print("✅ Answer generated successfully")
 
 class SupervisorNode(Node):
     def prep(self, shared):
         """Get the current answer for evaluation."""
         return shared["answer"]
     
-    def exec(self, answer):
+    def exec(self, prep_res):
         """Check if the answer is valid or nonsensical."""
-        print(f"    🔍 Supervisor checking answer quality...")
+        print("    🔍 Supervisor checking answer quality...")
         
         # Check for obvious markers of the nonsense answers
         nonsense_markers = [
@@ -149,7 +152,7 @@ def exec(self, answer):
         ]
         
         # Check if the answer contains any nonsense markers
-        is_nonsense = any(marker in answer for marker in nonsense_markers)
+        is_nonsense = any(marker in prep_res for marker in nonsense_markers)
         
         if is_nonsense:
             return {"valid": False, "reason": "Answer appears to be nonsensical or unhelpful"}
diff --git a/cookbook/pocketflow-tao/nodes.py b/cookbook/pocketflow-tao/nodes.py
index 261253a6..eb6276ad 100644
--- a/cookbook/pocketflow-tao/nodes.py
+++ b/cookbook/pocketflow-tao/nodes.py
@@ -1,9 +1,11 @@
 # nodes.py
 
-from pocketflow import Node
 import yaml
+
+from pocketflow import Node
 from utils import call_llm
 
+
 class ThinkNode(Node):
     def prep(self, shared):
         """Prepare the context needed for thinking"""
@@ -92,9 +94,9 @@ def prep(self, shared):
         action_input = shared["current_action_input"]
         return action, action_input
     
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Execute action and return result"""
-        action, action_input = inputs
+        action, action_input = prep_res
         
         print(f"🚀 Executing action: {action}, input: {action_input}")
         
@@ -118,7 +120,7 @@ def post(self, shared, prep_res, exec_res):
         """Save action result"""
         # Save the current action result
         shared["current_action_result"] = exec_res
-        print(f"✅ Action completed, result obtained")
+        print("✅ Action completed, result obtained")
         
         # Continue to observation node
         return "observe"
@@ -143,9 +145,9 @@ def prep(self, shared):
         action_result = shared["current_action_result"]
         return action, action_input, action_result
     
-    def exec(self, inputs):
+    def exec(self, prep_res):
         """Analyze action results, generate observation"""
-        action, action_input, action_result = inputs
+        action, action_input, action_result = prep_res
         
         # Build prompt
         prompt = f"""
diff --git a/cookbook/pocketflow-text2sql/nodes.py b/cookbook/pocketflow-text2sql/nodes.py
index 0842f2e0..23181a15 100644
--- a/cookbook/pocketflow-text2sql/nodes.py
+++ b/cookbook/pocketflow-text2sql/nodes.py
@@ -1,15 +1,18 @@
 import sqlite3
 import time
-import yaml # Import yaml here as nodes use it
+
+import yaml  # Import yaml here as nodes use it
+
 from pocketflow import Node
 from utils.call_llm import call_llm
 
+
 class GetSchema(Node):
     def prep(self, shared):
         return shared["db_path"]
 
-    def exec(self, db_path):
-        conn = sqlite3.connect(db_path)
+    def exec(self, prep_res):
+        conn = sqlite3.connect(prep_res)
         cursor = conn.cursor()
         cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
         tables = cursor.fetchall()
diff --git a/cookbook/pocketflow-tool-crawler/nodes.py b/cookbook/pocketflow-tool-crawler/nodes.py
index 81ce4a4a..5c75c08b 100644
--- a/cookbook/pocketflow-tool-crawler/nodes.py
+++ b/cookbook/pocketflow-tool-crawler/nodes.py
@@ -1,7 +1,8 @@
-from pocketflow import Node, BatchNode
 from tools.crawler import WebCrawler
 from tools.parser import analyze_site
-from typing import List, Dict
+
+from pocketflow import BatchNode, Node
+
 
 class CrawlWebsiteNode(Node):
     """Node to crawl a website and extract content"""
@@ -9,8 +10,8 @@ class CrawlWebsiteNode(Node):
     def prep(self, shared):
         return shared.get("base_url"), shared.get("max_pages", 10)
         
-    def exec(self, inputs):
-        base_url, max_pages = inputs
+    def exec(self, prep_res):
+        base_url, max_pages = prep_res
         if not base_url:
             return []
             
@@ -30,13 +31,13 @@ def prep(self, shared):
         batch_size = 5
         return [results[i:i+batch_size] for i in range(0, len(results), batch_size)]
         
-    def exec(self, batch):
-        return analyze_site(batch)
+    def exec(self, prep_res):
+        return analyze_site(prep_res)
         
-    def post(self, shared, prep_res, exec_res_list):
+    def post(self, shared, prep_res, exec_res):
         # Flatten results from all batches
         all_results = []
-        for batch_results in exec_res_list:
+        for batch_results in exec_res:
             all_results.extend(batch_results)
             
         shared["analyzed_results"] = all_results
@@ -48,15 +49,15 @@ class GenerateReportNode(Node):
     def prep(self, shared):
         return shared.get("analyzed_results", [])
         
-    def exec(self, results):
-        if not results:
+    def exec(self, prep_res):
+        if not prep_res:
             return "No results to report"
             
         report = []
-        report.append(f"Analysis Report\n")
-        report.append(f"Total pages analyzed: {len(results)}\n")
+        report.append("Analysis Report\n")
+        report.append(f"Total pages analyzed: {len(prep_res)}\n")
         
-        for page in results:
+        for page in prep_res:
             report.append(f"\nPage: {page['url']}")
             report.append(f"Title: {page['title']}")
             
diff --git a/cookbook/pocketflow-tool-database/nodes.py b/cookbook/pocketflow-tool-database/nodes.py
index 0f256f17..e1be72c7 100644
--- a/cookbook/pocketflow-tool-database/nodes.py
+++ b/cookbook/pocketflow-tool-database/nodes.py
@@ -1,10 +1,12 @@
-from pocketflow import Node
 from tools.database import execute_sql, init_db
 
+from pocketflow import Node
+
+
 class InitDatabaseNode(Node):
     """Node for initializing the database"""
     
-    def exec(self, _):
+    def exec(self, prep_res):
         init_db()
         return "Database initialized"
         
@@ -21,8 +23,8 @@ def prep(self, shared):
             shared.get("task_description", "")
         )
         
-    def exec(self, inputs):
-        title, description = inputs
+    def exec(self, prep_res):
+        title, description = prep_res
         query = "INSERT INTO tasks (title, description) VALUES (?, ?)"
         execute_sql(query, (title, description))
         return "Task created successfully"
@@ -34,7 +36,7 @@ def post(self, shared, prep_res, exec_res):
 class ListTasksNode(Node):
     """Node for listing all tasks"""
     
-    def exec(self, _):
+    def exec(self, prep_res):
         query = "SELECT * FROM tasks"
         return execute_sql(query)
         
diff --git a/cookbook/pocketflow-tool-embeddings/nodes.py b/cookbook/pocketflow-tool-embeddings/nodes.py
index ebf8e540..2c30b7b1 100644
--- a/cookbook/pocketflow-tool-embeddings/nodes.py
+++ b/cookbook/pocketflow-tool-embeddings/nodes.py
@@ -1,6 +1,8 @@
-from pocketflow import Node
 from tools.embeddings import get_embedding
 
+from pocketflow import Node
+
+
 class EmbeddingNode(Node):
     """Node for getting embeddings from OpenAI API"""
     
@@ -8,9 +10,9 @@ def prep(self, shared):
         # Get text from shared store
         return shared.get("text", "")
         
-    def exec(self, text):
+    def exec(self, prep_res):
         # Get embedding using tool function
-        return get_embedding(text)
+        return get_embedding(prep_res)
         
     def post(self, shared, prep_res, exec_res):
         # Store embedding in shared store
diff --git a/cookbook/pocketflow-tool-pdf-vision/nodes.py b/cookbook/pocketflow-tool-pdf-vision/nodes.py
index cc9250e6..14b29833 100644
--- a/cookbook/pocketflow-tool-pdf-vision/nodes.py
+++ b/cookbook/pocketflow-tool-pdf-vision/nodes.py
@@ -1,9 +1,11 @@
-from pocketflow import Node, BatchNode
+import os
+from pathlib import Path
+
 from tools.pdf import pdf_to_images
 from tools.vision import extract_text_from_image
-from typing import List, Dict, Any
-from pathlib import Path
-import os
+
+from pocketflow import BatchNode, Node
+
 
 class ProcessPDFBatchNode(BatchNode):
     """Node for processing multiple PDFs from a directory"""
@@ -30,25 +32,25 @@ def prep(self, shared):
         print(f"Found {len(pdf_files)} PDF files")
         return pdf_files
     
-    def exec(self, item):
+    def exec(self, prep_res):
         # Create flow for single PDF
         flow = create_single_pdf_flow()
         
         # Process PDF
-        print(f"\nProcessing: {os.path.basename(item['pdf_path'])}")
+        print(f"\nProcessing: {os.path.basename(prep_res['pdf_path'])}")
         print("-" * 50)
         
         # Run flow
-        shared = item.copy()
+        shared = prep_res.copy()
         flow.run(shared)
         
         return {
-            "filename": os.path.basename(item["pdf_path"]),
+            "filename": os.path.basename(prep_res["pdf_path"]),
             "text": shared.get("final_text", "No text extracted")
         }
     
-    def post(self, shared, prep_res, exec_res_list):
-        shared["results"] = exec_res_list
+    def post(self, shared, prep_res, exec_res):
+        shared["results"] = exec_res
         return "default"
 
 class LoadPDFNode(Node):
@@ -57,8 +59,8 @@ class LoadPDFNode(Node):
     def prep(self, shared):
         return shared.get("pdf_path", "")
         
-    def exec(self, pdf_path):
-        return pdf_to_images(pdf_path)
+    def exec(self, prep_res):
+        return pdf_to_images(prep_res)
         
     def post(self, shared, prep_res, exec_res):
         shared["page_images"] = exec_res
@@ -73,8 +75,8 @@ def prep(self, shared):
             shared.get("extraction_prompt", None)
         )
         
-    def exec(self, inputs):
-        images, prompt = inputs
+    def exec(self, prep_res):
+        images, prompt = prep_res
         results = []
         
         for img, page_num in images:
@@ -96,9 +98,9 @@ class CombineResultsNode(Node):
     def prep(self, shared):
         return shared.get("extracted_text", [])
         
-    def exec(self, results):
+    def exec(self, prep_res):
         # Sort by page number
-        sorted_results = sorted(results, key=lambda x: x["page"])
+        sorted_results = sorted(prep_res, key=lambda x: x["page"])
         
         # Combine text with page numbers
         combined = []
diff --git a/cookbook/pocketflow-tool-search/nodes.py b/cookbook/pocketflow-tool-search/nodes.py
index 4ecc9ea1..eb93ace9 100644
--- a/cookbook/pocketflow-tool-search/nodes.py
+++ b/cookbook/pocketflow-tool-search/nodes.py
@@ -1,7 +1,8 @@
-from pocketflow import Node
-from tools.search import SearchTool
 from tools.parser import analyze_results
-from typing import List, Dict
+from tools.search import SearchTool
+
+from pocketflow import Node
+
 
 class SearchNode(Node):
     """Node to perform web search using SerpAPI"""
@@ -9,8 +10,8 @@ class SearchNode(Node):
     def prep(self, shared):
         return shared.get("query"), shared.get("num_results", 5)
         
-    def exec(self, inputs):
-        query, num_results = inputs
+    def exec(self, prep_res):
+        query, num_results = prep_res
         if not query:
             return []
             
@@ -27,8 +28,8 @@ class AnalyzeResultsNode(Node):
     def prep(self, shared):
         return shared.get("query"), shared.get("search_results", [])
         
-    def exec(self, inputs):
-        query, results = inputs
+    def exec(self, prep_res):
+        query, results = prep_res
         if not results:
             return {
                 "summary": "No search results to analyze",
diff --git a/cookbook/pocketflow-tracing/README.md b/cookbook/pocketflow-tracing/README.md
index e5c4a797..1bea074b 100644
--- a/cookbook/pocketflow-tracing/README.md
+++ b/cookbook/pocketflow-tracing/README.md
@@ -49,8 +49,8 @@ class MyNode(Node):
     def prep(self, shared):
         return shared["input"]
     
-    def exec(self, data):
-        return f"Processed: {data}"
+    def exec(self, prep_res):
+        return f"Processed: {prep_res}"
     
     def post(self, shared, prep_res, exec_res):
         shared["output"] = exec_res
diff --git a/cookbook/pocketflow-tracing/examples/async_example.py b/cookbook/pocketflow-tracing/examples/async_example.py
index 5d269629..925c48d2 100644
--- a/cookbook/pocketflow-tracing/examples/async_example.py
+++ b/cookbook/pocketflow-tracing/examples/async_example.py
@@ -7,8 +7,9 @@
 """
 
 import asyncio
-import sys
 import os
+import sys
+
 from dotenv import load_dotenv
 
 # Load environment variables
@@ -18,8 +19,9 @@
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
 sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
 
-from pocketflow import AsyncNode, AsyncFlow
-from tracing import trace_flow, TracingConfig
+from tracing import trace_flow
+
+from pocketflow import AsyncFlow, AsyncNode
 
 
 class AsyncDataFetchNode(AsyncNode):
@@ -30,17 +32,17 @@ async def prep_async(self, shared):
         query = shared.get("query", "default")
         return query
 
-    async def exec_async(self, query):
+    async def exec_async(self, prep_res):
         """Simulate async data fetching."""
-        print(f"🔍 Fetching data for query: {query}")
+        print(f"🔍 Fetching data for query: {prep_res}")
 
         # Simulate async operation
         await asyncio.sleep(1)
 
         # Return mock data
         data = {
-            "query": query,
-            "results": [f"Result {i} for {query}" for i in range(3)],
+            "query": prep_res,
+            "results": [f"Result {i} for {prep_res}" for i in range(3)],
             "timestamp": "2024-01-01T00:00:00Z",
         }
         return data
@@ -58,7 +60,7 @@ async def prep_async(self, shared):
         """Get the fetched data."""
         return shared.get("fetched_data", {})
 
-    async def exec_async(self, data):
+    async def exec_async(self, prep_res):
         """Process the data asynchronously."""
         print("⚙️ Processing fetched data...")
 
@@ -67,11 +69,11 @@ async def exec_async(self, data):
 
         # Process the results
         processed_results = []
-        for result in data.get("results", []):
+        for result in prep_res.get("results", []):
             processed_results.append(f"PROCESSED: {result}")
 
         return {
-            "original_query": data.get("query"),
+            "original_query": prep_res.get("query"),
             "processed_results": processed_results,
             "result_count": len(processed_results),
         }
diff --git a/cookbook/pocketflow-tracing/examples/basic_example.py b/cookbook/pocketflow-tracing/examples/basic_example.py
index 587f76d4..23b4fd73 100644
--- a/cookbook/pocketflow-tracing/examples/basic_example.py
+++ b/cookbook/pocketflow-tracing/examples/basic_example.py
@@ -6,8 +6,9 @@
 trace a simple PocketFlow workflow.
 """
 
-import sys
 import os
+import sys
+
 from dotenv import load_dotenv
 
 # Load environment variables
@@ -17,8 +18,9 @@
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
 sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
 
-from pocketflow import Node, Flow
-from tracing import trace_flow, TracingConfig
+from tracing import trace_flow
+
+from pocketflow import Flow, Node
 
 
 class GreetingNode(Node):
@@ -29,9 +31,9 @@ def prep(self, shared):
         name = shared.get("name", "World")
         return name
 
-    def exec(self, name):
+    def exec(self, prep_res):
         """Create a greeting message."""
-        greeting = f"Hello, {name}!"
+        greeting = f"Hello, {prep_res}!"
         return greeting
 
     def post(self, shared, prep_res, exec_res):
@@ -47,9 +49,9 @@ def prep(self, shared):
         """Get the greeting from shared data."""
         return shared.get("greeting", "")
 
-    def exec(self, greeting):
+    def exec(self, prep_res):
         """Convert to uppercase."""
-        return greeting.upper()
+        return prep_res.upper()
 
     def post(self, shared, prep_res, exec_res):
         """Store the uppercase greeting."""
diff --git a/cookbook/pocketflow-voice-chat/nodes.py b/cookbook/pocketflow-voice-chat/nodes.py
index a43ba382..54ca9a02 100644
--- a/cookbook/pocketflow-voice-chat/nodes.py
+++ b/cookbook/pocketflow-voice-chat/nodes.py
@@ -1,17 +1,18 @@
-import numpy as np
-import scipy.io.wavfile
 import io
-import soundfile # For converting MP3 bytes to NumPy array
+
+import scipy.io.wavfile
+import soundfile  # For converting MP3 bytes to NumPy array
 
 from pocketflow import Node
-from utils.audio_utils import record_audio, play_audio_data
-from utils.speech_to_text import speech_to_text_api
+from utils.audio_utils import play_audio_data, record_audio
 from utils.call_llm import call_llm
+from utils.speech_to_text import speech_to_text_api
 from utils.text_to_speech import text_to_speech_api
 
+
 class CaptureAudioNode(Node):
     """Records audio input from the user using VAD."""
-    def exec(self, _): # prep_res is not used as per design
+    def exec(self, prep_res): # prep_res is not used as per design
         print("\nListening for your query...")
         audio_data, sample_rate = record_audio()
         if audio_data is None:
diff --git a/cookbook/pocketflow-workflow/nodes.py b/cookbook/pocketflow-workflow/nodes.py
index 00035f87..1e10b6ac 100644
--- a/cookbook/pocketflow-workflow/nodes.py
+++ b/cookbook/pocketflow-workflow/nodes.py
@@ -1,15 +1,16 @@
-import re
-from pocketflow import Node, BatchNode
-from utils.call_llm import call_llm
 import yaml
 
+from pocketflow import BatchNode, Node
+from utils.call_llm import call_llm
+
+
 class GenerateOutline(Node):
     def prep(self, shared):
         return shared["topic"]
     
-    def exec(self, topic):
+    def exec(self, prep_res):
         prompt = f"""
-Create a simple outline for an article about {topic}.
+Create a simple outline for an article about {prep_res}.
 Include at most 3 main sections (no subsections).
 
 Output the sections in YAML format as shown below:
@@ -55,11 +56,11 @@ def prep(self, shared):
         self.sections = shared.get("sections", [])
         return self.sections
     
-    def exec(self, section):
+    def exec(self, prep_res):
         prompt = f"""
 Write a short paragraph (MAXIMUM 100 WORDS) about this section:
 
-{section}
+{prep_res}
 
 Requirements:
 - Explain the idea in simple, easy-to-understand terms
@@ -70,18 +71,18 @@ def exec(self, section):
         content = call_llm(prompt)
         
         # Show progress for this section
-        current_section_index = self.sections.index(section) if section in self.sections else 0
+        current_section_index = self.sections.index(prep_res) if prep_res in self.sections else 0
         total_sections = len(self.sections)
-        print(f"✓ Completed section {current_section_index + 1}/{total_sections}: {section}")
+        print(f"✓ Completed section {current_section_index + 1}/{total_sections}: {prep_res}")
         
-        return section, content
+        return prep_res, content
     
-    def post(self, shared, prep_res, exec_res_list):
-        # exec_res_list contains [(section, content), (section, content), ...]
+    def post(self, shared, prep_res, exec_res):
+        # exec_res contains [(section, content), (section, content), ...]
         section_contents = {}
         all_sections_content = []
         
-        for section, content in exec_res_list:
+        for section, content in exec_res:
             section_contents[section] = content
             all_sections_content.append(f"## {section}\n\n{content}\n")
         
@@ -107,14 +108,14 @@ def prep(self, shared):
         """
         return shared["draft"]
     
-    def exec(self, draft):
+    def exec(self, prep_res):
         """
         Apply a specific style to the article
         """
         prompt = f"""
         Rewrite the following draft in a conversational, engaging style:
         
-        {draft}
+        {prep_res}
         
         Make it:
         - Conversational and warm in tone
diff --git a/cookbook/pocketflow_demo.ipynb b/cookbook/pocketflow_demo.ipynb
index 8c8a5761..7c4e0788 100644
--- a/cookbook/pocketflow_demo.ipynb
+++ b/cookbook/pocketflow_demo.ipynb
@@ -150,7 +150,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -169,9 +169,9 @@
     "        # Read data from shared store\n",
     "        return shared[\"data\"][\"before.txt\"]\n",
     "        \n",
-    "    def exec(self, text):\n",
+    "    def exec(self, prep_res):\n",
     "        # Call LLM to summarize\n",
-    "        prompt = f\"Summarize this text in 50 words:\\n\\n{text}\"\n",
+    "        prompt = f\"Summarize this text in 50 words:\\n\\n{prep_res}\"\n",
     "        return call_llm(prompt)\n",
     "    \n",
     "    def post(self, shared, prep_res, exec_res):\n",
@@ -252,7 +252,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -287,19 +287,19 @@
     "        # Return list of (filename, content) tuples from shared store\n",
     "        return [(fn, content) for fn, content in shared[\"data\"].items()]\n",
     "        \n",
-    "    def exec(self, item):\n",
+    "    def exec(self, prep_res):\n",
     "        # Unpack the filename and content\n",
-    "        filename, text = item\n",
+    "        filename, text = prep_res\n",
     "        # Call LLM to summarize\n",
     "        prompt = f\"Summarize this text in 50 words:\\n\\n{text}\"\n",
     "        summary = call_llm(prompt)\n",
     "        return filename, summary\n",
     "    \n",
-    "    def post(self, shared, prep_res, exec_res_list):\n",
+    "    def post(self, shared, prep_res, exec_res):\n",
     "        # Store all summaries in a dict by filename\n",
     "        shared[\"summaries\"] = {\n",
     "            filename: summary \n",
-    "            for filename, summary in exec_res_list\n",
+    "            for filename, summary in exec_res\n",
     "        }\n",
     "        return \"default\"\n",
     "\n",
@@ -397,7 +397,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -437,11 +437,11 @@
     "        # Get list of (filename, content) pairs\n",
     "        return list(shared[\"data\"].items())\n",
     "        \n",
-    "    def exec(self, items):\n",
+    "    def exec(self, prep_res):\n",
     "        # Create embeddings for each document\n",
     "        embeddings = []\n",
     "        filenames = []\n",
-    "        for filename, content in items:\n",
+    "        for filename, content in prep_res:\n",
     "            embedding = get_embedding(content)\n",
     "            embeddings.append(embedding)\n",
     "            filenames.append(filename)\n",
@@ -468,12 +468,12 @@
     "            return None\n",
     "        return question\n",
     "        \n",
-    "    def exec(self, question):\n",
-    "        if question is None:\n",
+    "    def exec(self, prep_res):\n",
+    "        if prep_res is None:\n",
     "            return None\n",
     "            \n",
     "        # Get question embedding and search\n",
-    "        query_embedding = get_embedding(question)\n",
+    "        query_embedding = get_embedding(prep_res)\n",
     "        \n",
     "        # Search for most similar document\n",
     "        D, I = shared[\"search_index\"].search(\n",
@@ -483,7 +483,7 @@
     "        most_relevant_idx = I[0][0]\n",
     "        most_relevant_file = shared[\"filenames\"][most_relevant_idx]\n",
     "        \n",
-    "        return question, most_relevant_file\n",
+    "        return prep_res, most_relevant_file\n",
     "        \n",
     "    def post(self, shared, prep_res, exec_res):\n",
     "        if exec_res is None:\n",
@@ -502,8 +502,8 @@
     "            shared[\"context\"]\n",
     "        )\n",
     "        \n",
-    "    def exec(self, inputs):\n",
-    "        question, context = inputs\n",
+    "    def exec(self, prep_res):\n",
+    "        question, context = prep_res\n",
     "        prompt = f\"\"\"\n",
     "Context: {context}\n",
     "\n",
diff --git a/docs/core_abstraction/batch.md b/docs/core_abstraction/batch.md
index 35063bda..343ddc68 100644
--- a/docs/core_abstraction/batch.md
+++ b/docs/core_abstraction/batch.md
@@ -31,8 +31,8 @@ class MapSummaries(BatchNode):
         chunks = [content[i:i+chunk_size] for i in range(0, len(content), chunk_size)]
         return chunks
 
-    def exec(self, chunk):
-        prompt = f"Summarize this chunk in 10 words: {chunk}"
+    def exec(self, prep_res):
+        prompt = f"Summarize this chunk in 10 words: {prep_res}"
         summary = call_llm(prompt)
         return summary
 
@@ -77,8 +77,8 @@ class LoadFile(Node):
         filename = self.params["filename"]  # Important! Use self.params, not shared
         return filename
         
-    def exec(self, filename):
-        with open(filename, 'r') as f:
+    def exec(self, prep_res):
+        with open(prep_res, 'r') as f:
             return f.read()
             
     def post(self, shared, prep_res, exec_res):
@@ -91,8 +91,8 @@ class Summarize(Node):
     def prep(self, shared):
         return shared["current_file_content"]
         
-    def exec(self, content):
-        prompt = f"Summarize this file in 50 words: {content}"
+    def exec(self, prep_res):
+        prompt = f"Summarize this file in 50 words: {prep_res}"
         return call_llm(prompt)
         
     def post(self, shared, prep_res, exec_res):
@@ -156,9 +156,9 @@ class ProcessFile(Node):
         full_path = os.path.join(directory, filename)
         return full_path
         
-    def exec(self, full_path):
+    def exec(self, prep_res):
         # Process the file...
-        return f"Processed {full_path}"
+        return f"Processed {prep_res}"
         
     def post(self, shared, prep_res, exec_res):
         # Store results, perhaps indexed by path
diff --git a/docs/design_pattern/agent.md b/docs/design_pattern/agent.md
index a0864264..08866020 100644
--- a/docs/design_pattern/agent.md
+++ b/docs/design_pattern/agent.md
@@ -80,8 +80,8 @@ class DecideAction(Node):
         query = shared["query"]
         return query, context
         
-    def exec(self, inputs):
-        query, context = inputs
+    def exec(self, prep_res):
+        query, context = prep_res
         prompt = f"""
 Given input: {query}
 Previous search results: {context}
@@ -114,8 +114,8 @@ class SearchWeb(Node):
     def prep(self, shared):
         return shared["search_term"]
         
-    def exec(self, search_term):
-        return search_web(search_term)
+    def exec(self, prep_res):
+        return search_web(prep_res)
     
     def post(self, shared, prep_res, exec_res):
         prev_searches = shared.get("context", [])
@@ -128,8 +128,8 @@ class DirectAnswer(Node):
     def prep(self, shared):
         return shared["query"], shared.get("context", "")
         
-    def exec(self, inputs):
-        query, context = inputs
+    def exec(self, prep_res):
+        query, context = prep_res
         return call_llm(f"Context: {context}\nAnswer: {query}")
 
     def post(self, shared, prep_res, exec_res):
diff --git a/docs/design_pattern/mapreduce.md b/docs/design_pattern/mapreduce.md
index 237b3cb9..09464059 100644
--- a/docs/design_pattern/mapreduce.md
+++ b/docs/design_pattern/mapreduce.md
@@ -27,8 +27,8 @@ class SummarizeAllFiles(BatchNode):
         files_dict = shared["files"]  # e.g. 10 files
         return list(files_dict.items())  # [("file1.txt", "aaa..."), ("file2.txt", "bbb..."), ...]
 
-    def exec(self, one_file):
-        filename, file_content = one_file
+    def exec(self, prep_res):
+        filename, file_content = prep_res
         summary_text = call_llm(f"Summarize the following file:\n{file_content}")
         return (filename, summary_text)
 
@@ -39,17 +39,17 @@ class CombineSummaries(Node):
     def prep(self, shared):
         return shared["file_summaries"]
 
-    def exec(self, file_summaries):
+    def exec(self, prep_res):
         # format as: "File1: summary\nFile2: summary...\n"
         text_list = []
-        for fname, summ in file_summaries.items():
+        for fname, summ in prep_res.items():
             text_list.append(f"{fname} summary:\n{summ}\n")
         big_text = "\n---\n".join(text_list)
 
         return call_llm(f"Combine these file summaries into one final summary:\n{big_text}")
 
-    def post(self, shared, prep_res, final_summary):
-        shared["all_files_summary"] = final_summary
+    def post(self, shared, prep_res, exec_res):
+        shared["all_files_summary"] = exec_res
 
 batch_node = SummarizeAllFiles()
 combine_node = CombineSummaries()
diff --git a/docs/design_pattern/multi_agent.md b/docs/design_pattern/multi_agent.md
index f0a7bcb5..da9c27e4 100644
--- a/docs/design_pattern/multi_agent.md
+++ b/docs/design_pattern/multi_agent.md
@@ -83,10 +83,10 @@ class AsyncHinter(AsyncNode):
             return None
         return shared["target_word"], shared["forbidden_words"], shared.get("past_guesses", [])
 
-    async def exec_async(self, inputs):
-        if inputs is None:
+    async def exec_async(self, prep_res):
+        if prep_res is None:
             return None
-        target, forbidden, past_guesses = inputs
+        target, forbidden, past_guesses = prep_res
         prompt = f"Generate hint for '{target}'\nForbidden words: {forbidden}"
         if past_guesses:
             prompt += f"\nPrevious wrong guesses: {past_guesses}\nMake hint more specific."
@@ -107,8 +107,8 @@ class AsyncGuesser(AsyncNode):
         hint = await shared["guesser_queue"].get()
         return hint, shared.get("past_guesses", [])
 
-    async def exec_async(self, inputs):
-        hint, past_guesses = inputs
+    async def exec_async(self, prep_res):
+        hint, past_guesses = prep_res
         prompt = f"Given hint: {hint}, past wrong guesses: {past_guesses}, make a new guess. Directly reply a single word:"
         guess = call_llm(prompt)
         print(f"Guesser: I guess it's - {guess}")
diff --git a/docs/design_pattern/rag.md b/docs/design_pattern/rag.md
index a5347825..cf41b7de 100644
--- a/docs/design_pattern/rag.md
+++ b/docs/design_pattern/rag.md
@@ -30,9 +30,9 @@ class ChunkDocs(BatchNode):
         # A list of file paths in shared["files"]. We process each file.
         return shared["files"]
 
-    def exec(self, filepath):
+    def exec(self, prep_res):
         # read file content. In real usage, do error handling.
-        with open(filepath, "r", encoding="utf-8") as f:
+        with open(prep_res, "r", encoding="utf-8") as f:
             text = f.read()
         # chunk by 100 chars each
         chunks = []
@@ -53,8 +53,8 @@ class EmbedDocs(BatchNode):
     def prep(self, shared):
         return shared["all_chunks"]
 
-    def exec(self, chunk):
-        return get_embedding(chunk)
+    def exec(self, prep_res):
+        return get_embedding(prep_res)
 
     def post(self, shared, prep_res, exec_res_list):
         # Store the list of embeddings.
@@ -66,13 +66,13 @@ class StoreIndex(Node):
         # We'll read all embeds from shared.
         return shared["all_embeds"]
 
-    def exec(self, all_embeds):
+    def exec(self, prep_res):
         # Create a vector index (faiss or other DB in real usage).
-        index = create_index(all_embeds)
+        index = create_index(prep_res)
         return index
 
-    def post(self, shared, prep_res, index):
-        shared["index"] = index
+    def post(self, shared, prep_res, exec_res):
+        shared["index"] = exec_res
 
 # Wire them in sequence
 chunk_node = ChunkDocs()
@@ -106,40 +106,40 @@ class EmbedQuery(Node):
     def prep(self, shared):
         return shared["question"]
 
-    def exec(self, question):
-        return get_embedding(question)
+    def exec(self, prep_res):
+        return get_embedding(prep_res)
 
-    def post(self, shared, prep_res, q_emb):
-        shared["q_emb"] = q_emb
+    def post(self, shared, prep_res, exec_res):
+        shared["q_emb"] = exec_res
 
 class RetrieveDocs(Node):
     def prep(self, shared):
         # We'll need the query embedding, plus the offline index/chunks
         return shared["q_emb"], shared["index"], shared["all_chunks"]
 
-    def exec(self, inputs):
-        q_emb, index, chunks = inputs
+    def exec(self, prep_res):
+        q_emb, index, chunks = prep_res
         I, D = search_index(index, q_emb, top_k=1)
         best_id = I[0][0]
         relevant_chunk = chunks[best_id]
         return relevant_chunk
 
-    def post(self, shared, prep_res, relevant_chunk):
-        shared["retrieved_chunk"] = relevant_chunk
-        print("Retrieved chunk:", relevant_chunk[:60], "...")
+    def post(self, shared, prep_res, exec_res):
+        shared["retrieved_chunk"] = exec_res
+        print("Retrieved chunk:", exec_res[:60], "...")
 
 class GenerateAnswer(Node):
     def prep(self, shared):
         return shared["question"], shared["retrieved_chunk"]
 
-    def exec(self, inputs):
-        question, chunk = inputs
+    def exec(self, prep_res):
+        question, chunk = prep_res
         prompt = f"Question: {question}\nContext: {chunk}\nAnswer:"
         return call_llm(prompt)
 
-    def post(self, shared, prep_res, answer):
-        shared["answer"] = answer
-        print("Answer:", answer)
+    def post(self, shared, prep_res, exec_res):
+        shared["answer"] = exec_res
+        print("Answer:", exec_res)
 
 embed_qnode = EmbedQuery()
 retrieve_node = RetrieveDocs()
diff --git a/docs/design_pattern/workflow.md b/docs/design_pattern/workflow.md
index 92dc536e..5bbe0ac5 100644
--- a/docs/design_pattern/workflow.md
+++ b/docs/design_pattern/workflow.md
@@ -24,17 +24,17 @@ Many real-world tasks are too complex for one LLM call. The solution is to **Tas
 ```python
 class GenerateOutline(Node):
     def prep(self, shared): return shared["topic"]
-    def exec(self, topic): return call_llm(f"Create a detailed outline for an article about {topic}")
+    def exec(self, prep_res): return call_llm(f"Create a detailed outline for an article about {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["outline"] = exec_res
 
 class WriteSection(Node):
     def prep(self, shared): return shared["outline"]
-    def exec(self, outline): return call_llm(f"Write content based on this outline: {outline}")
+    def exec(self, prep_res): return call_llm(f"Write content based on this outline: {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["draft"] = exec_res
 
 class ReviewAndRefine(Node):
     def prep(self, shared): return shared["draft"]
-    def exec(self, draft): return call_llm(f"Review and improve this draft: {draft}")
+    def exec(self, prep_res): return call_llm(f"Review and improve this draft: {prep_res}")
     def post(self, shared, prep_res, exec_res): shared["final_article"] = exec_res
 
 # Connect nodes
diff --git a/docs/guide.md b/docs/guide.md
index 1b45bdcd..c5afe5cd 100644
--- a/docs/guide.md
+++ b/docs/guide.md
@@ -167,7 +167,7 @@ my_project/
   from utils.call_llm import call_llm
 
   class GetQuestionNode(Node):
-      def exec(self, _):
+      def exec(self, prep_res):
           # Get question directly from user input
           user_question = input("Enter your question: ")
           return user_question
@@ -182,9 +182,9 @@ my_project/
           # Read question from shared
           return shared["question"]
       
-      def exec(self, question):
+      def exec(self, prep_res):
           # Call LLM to get the answer
-          return call_llm(question)
+          return call_llm(prep_res)
       
       def post(self, shared, prep_res, exec_res):
           # Store the answer in shared
diff --git a/docs/utility_function/llm.md b/docs/utility_function/llm.md
index a849f79d..a3216a8a 100644
--- a/docs/utility_function/llm.md
+++ b/docs/utility_function/llm.md
@@ -141,8 +141,8 @@ def call_llm(prompt, use_cache):
     return cached_call.__wrapped__(prompt)
 
 class SummarizeNode(Node):
-    def exec(self, text):
-        return call_llm(f"Summarize: {text}", self.cur_retry==0)
+    def exec(self, prep_res):
+        return call_llm(f"Summarize: {prep_res}", self.cur_retry==0)
 ```
 
 - Enable logging:
diff --git a/pocketflow/__init__.py b/pocketflow/__init__.py
index a7203dfe..7bbf47be 100644
--- a/pocketflow/__init__.py
+++ b/pocketflow/__init__.py
@@ -1,100 +1,810 @@
-import asyncio, warnings, copy, time
+from __future__ import annotations
+
+import asyncio
+import copy
+import time
+import warnings
+from typing import Any
+
+type ParamValue = str | int | float | bool | None | list[Any] | dict[str, Any]
+
 
 class BaseNode:
-    def __init__(self): self.params,self.successors={},{}
-    def set_params(self,params): self.params=params
-    def next(self,node,action="default"):
-        if action in self.successors: warnings.warn(f"Overwriting successor for action '{action}'")
-        self.successors[action]=node; return node
-    def prep(self,shared): pass
-    def exec(self,prep_res): pass
-    def post(self,shared,prep_res,exec_res): pass
-    def _exec(self,prep_res): return self.exec(prep_res)
-    def _run(self,shared): p=self.prep(shared); e=self._exec(p); return self.post(shared,p,e)
-    def run(self,shared): 
-        if self.successors: warnings.warn("Node won't run successors. Use Flow.")  
-        return self._run(shared)
-    def __rshift__(self,other): return self.next(other)
-    def __sub__(self,action):
-        if isinstance(action,str): return _ConditionalTransition(self,action)
+    """The base class for all nodes and flows in PocketFlow."""
+
+    def __init__(self) -> None:
+        """Initializes a `BaseNode` object."""
+        self.params: dict[str, ParamValue] = {}
+        self.successors: dict[str, "BaseNode"] = {}
+
+    def set_params(self, params: dict[str, ParamValue]) -> None:
+        """Sets the parameters of the node.
+
+        Args:
+            params (dict[str, ParamValue]): The parameters to set.
+        """
+        self.params = params
+
+    def next(self, node: "BaseNode", action: str = "default") -> "BaseNode":
+        """Sets the next node in the flow.
+
+        Args:
+            node (BaseNode): The next node in the flow.
+            action (str, optional): The action that triggers the transition to
+                the next node. Defaults to "default".
+
+        Returns:
+            BaseNode: The next node in the flow.
+        """
+        if action in self.successors:
+            warnings.warn(message=f"Overwriting successor for action '{action}'")
+        self.successors[action] = node
+        return node
+
+    def prep(self, shared: dict[str, Any]) -> Any | None:
+        """The pre-processing logic of the node.
+
+        This method is called before the node is executed.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any | None: The result of the pre-processing logic.
+        """
+        pass
+
+    def exec(self, prep_res: Any) -> Any | None:
+        """The execution logic of the node.
+
+        This method is called after the `prep` method.
+
+        Args:
+            prep_res (Any): The result of the `prep` method.
+
+        Returns:
+            Any | None: The result of the execution logic.
+        """
+        pass
+
+    def post(
+        self, shared: dict[str, Any], prep_res: Any, exec_res: Any
+    ) -> Any | None:
+        """The post-processing logic of the node.
+
+        This method is called after the node is executed.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+            prep_res (Any): The result of the `prep` method.
+            exec_res (Any): The result of the `exec` method.
+
+        Returns:
+            Any | None: The result of the post-processing logic.
+        """
+        pass
+
+    def _exec(self, prep_res: Any) -> Any:
+        """The execution wrapper for the node.
+
+        This method wraps the `exec` method.
+
+        Args:
+            prep_res (Any): The result of the `prep` method.
+
+        Returns:
+            Any: The result of the `exec` method.
+        """
+        return self.exec(prep_res=prep_res)
+
+    def _run(self, shared: dict[str, Any]) -> Any:
+        """The execution wrapper for the node.
+
+        This method wraps the `prep`, `_exec`, and `post` methods.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post` method.
+        """
+        p = self.prep(shared=shared)
+        e = self._exec(prep_res=p)
+        return self.post(shared=shared, prep_res=p, exec_res=e)
+
+    def run(self, shared: dict[str, Any]) -> Any:
+        """Runs the node.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post` method.
+        """
+        if self.successors:
+            warnings.warn(message="Node won't run successors. Use Flow.")
+        return self._run(shared=shared)
+
+    def __rshift__(self, other: "BaseNode") -> "BaseNode":
+        """Sets the next node in the flow.
+
+        This method is an alias for the `next` method.
+
+        Args:
+            other (BaseNode): The next node in the flow.
+
+        Returns:
+            BaseNode: The next node in the flow.
+        """
+        return self.next(node=other)
+
+    def __sub__(self, action: str) -> _ConditionalTransition:
+        """Creates a conditional transition.
+
+        Args:
+            action (str): The action that triggers the transition.
+
+        Raises:
+            TypeError: If the action is not a string.
+
+        Returns:
+            _ConditionalTransition: The conditional transition.
+        """
+        if isinstance(action, str):
+            return _ConditionalTransition(src=self, action=action)
         raise TypeError("Action must be a string")
 
+
 class _ConditionalTransition:
-    def __init__(self,src,action): self.src,self.action=src,action
-    def __rshift__(self,tgt): return self.src.next(tgt,self.action)
+    """A conditional transition between two nodes."""
+
+    def __init__(self, src: BaseNode, action: str) -> None:
+        """Initializes a `_ConditionalTransition` object.
+
+        Args:
+            src (BaseNode): The source node.
+            action (str): The action that triggers the transition.
+        """
+        self.src, self.action = src, action
+
+    def __rshift__(self, tgt: BaseNode) -> BaseNode:
+        """Sets the target node of the transition.
+
+        Args:
+            tgt (BaseNode): The target node.
+
+        Returns:
+            BaseNode: The source node.
+        """
+        return self.src.next(node=tgt, action=self.action)
+
 
 class Node(BaseNode):
-    def __init__(self,max_retries=1,wait=0): super().__init__(); self.max_retries,self.wait=max_retries,wait
-    def exec_fallback(self,prep_res,exc): raise exc
-    def _exec(self,prep_res):
+    """A `Node` object is an atomic, executable unit of a `Flow`.
+
+    It is characterized by the following:
+    - A single entry point and a single exit point.
+    - The ability to be retried on failure.
+    - A fallback mechanism for when all retries are exhausted.
+
+    Example:
+
+    .. code-block:: python
+        
+        class MyNode(Node):
+            def prep(self, shared: dict[str, Any]) -> Any:
+                # Prepare some data, often by looking at the shared storage
+                return {"data": shared.get("input", "default")}
+            def exec(self, prep_res: Any) -> Any:
+                # Execute some logic
+                return prep_res["data"].upper()
+            def post(self, shared: dict[str, Any], prep_res: Any, exec_res: Any) -> Any:
+                # Post-process the result
+                shared["result"] = exec_res
+                return "action_name"  # This will determine the next node in the flow
+                # returning None is the same as returning "default"
+    
+    """
+
+    def __init__(self, max_retries: int = 1, wait: int = 0) -> None:
+        """Initializes a `Node` object.
+
+        Args:
+            max_retries (int, optional): The maximum number of times to retry the
+                node on failure. Defaults to 1.
+            wait (int, optional): The number of seconds to wait between retries.
+                Defaults to 0.
+        """
+        super().__init__()
+        self.max_retries, self.wait = max_retries, wait
+        self.cur_retry: int = 0
+
+    def exec_fallback(self, prep_res: Any, exc: Exception) -> Any:
+        """The fallback mechanism for the node.
+
+        This method is called when the node has been retried `max_retries` times
+        and has failed every time.
+
+        Args:
+            prep_res (Any): The result of the `prep` method.
+            exc (Exception): The exception that was raised.
+
+        Raises:
+            exc: The exception that was raised.
+        """
+        raise exc
+
+    def _exec(self, prep_res: Any) -> Any:
+        """The execution wrapper for the node.
+
+        This method wraps the `exec` method with retry and fallback logic.
+
+        Args:
+            prep_res (Any): The result of the `prep` method.
+
+        Returns:
+            Any: The result of the `exec` method or the `exec_fallback` method.
+        """
         for self.cur_retry in range(self.max_retries):
-            try: return self.exec(prep_res)
+            try:
+                return self.exec(prep_res=prep_res)
             except Exception as e:
-                if self.cur_retry==self.max_retries-1: return self.exec_fallback(prep_res,e)
-                if self.wait>0: time.sleep(self.wait)
+                if self.cur_retry == self.max_retries - 1:
+                    return self.exec_fallback(prep_res=prep_res, exc=e)
+                if self.wait > 0:
+                    time.sleep(self.wait)
+        return None  # Should be unreachable
+
 
 class BatchNode(Node):
-    def _exec(self,items): return [super(BatchNode,self)._exec(i) for i in (items or [])]
+    """A `BatchNode` object is a `Node` that processes a batch of items.
+
+    The prep method should return a list of items to be processed
+    The exec method should process a single item and return a single result
+    The post method receives the list of results amd can return an action
+    """
+
+    def _exec(self, prep_res: list[Any] | None) -> list[Any]:
+        """The execution wrapper for the batch node.
+
+        This method wraps the `_exec` method of the parent class to process a
+        batch of items.
+
+        Args:
+            prep_res (list[Any] | None): The result of the `prep` method.
+
+        Returns:
+            list[Any]: The list of results.
+        """
+        return [
+            super(BatchNode, self)._exec(prep_res=i) for i in (prep_res or [])
+        ]
+
 
 class Flow(BaseNode):
-    def __init__(self,start=None): super().__init__(); self.start_node=start
-    def start(self,start): self.start_node=start; return start
-    def get_next_node(self,curr,action):
-        nxt=curr.successors.get(action or "default")
-        if not nxt and curr.successors: warnings.warn(f"Flow ends: '{action}' not found in {list(curr.successors)}")
+    """A `Flow` object orchestrates a graph of `Node` objects.
+
+    It is a state machine that transitions between nodes based on the `action`
+    returned by each node's `post` method.
+
+    Example:
+
+    .. code-block:: python
+        
+        # Create a flow
+        node1 = MyNode1()
+        node2 = MyNode2()
+        node1 >> node2  # Connects node1 to node2 with the default action
+        flow = Flow(start=node1)
+        ## the result of the last post method
+
+        # prepare input data
+        shared={"input": "hello"}
+
+        # run the flow and get the result of the flow's post method
+        # which is normally the result of the last node's post method
+        result = flow.run(shared=shared)
+        ## or examine an updated shared storage
+        print(shared)
+    """
+
+    def __init__(self, start: BaseNode | None = None) -> None:
+        """Initializes a `Flow` object.
+
+        Args:
+            start (BaseNode | None, optional): The start node of the flow.
+                Defaults to None.
+        """
+        super().__init__()
+        self.start_node = start
+
+    def start(self, start: BaseNode) -> BaseNode:
+        """Sets the start node of the flow.
+
+        Args:
+            start (BaseNode): The start node of the flow.
+
+        Returns:
+            BaseNode: The start node of the flow.
+        """
+        self.start_node = start
+        return start
+
+    def get_next_node(self, curr: BaseNode, action: str | None) -> BaseNode | None:
+        """Gets the next node in the flow.
+
+        Args:
+            curr (BaseNode): The current node.
+            action (str | None): The action returned by the current node.
+
+        Returns:
+            BaseNode | None: The next node in the flow, or `None` if the flow
+                has ended.
+        """
+        nxt = curr.successors.get(action or "default")
+        if not nxt and curr.successors:
+            warnings.warn(
+                message=f"Flow ends: '{action}' not found in {list(curr.successors)}"
+            )
         return nxt
-    def _orch(self,shared,params=None):
-        curr,p,last_action =copy.copy(self.start_node),(params or {**self.params}),None
-        while curr: curr.set_params(p); last_action=curr._run(shared); curr=copy.copy(self.get_next_node(curr,last_action))
+
+    def _orch(
+        self,
+        shared: dict[str, Any],
+        params: dict[str, ParamValue] | None = None,
+    ) -> str | None:
+        """The orchestration logic of the flow.
+
+        This method executes the nodes in the flow in the order they are
+        connected.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+            params (dict[str, ParamValue] | None, optional): The parameters to
+                pass to the nodes. Defaults to None.
+
+        Returns:
+            str | None: The last action returned by a node in the flow.
+        """
+        curr, p, last_action = (
+            copy.copy(self.start_node),
+            (params or {**self.params}),
+            None,
+        )
+        while curr:
+            curr.set_params(params=p)
+            last_action = curr._run(shared=shared)
+            curr = copy.copy(self.get_next_node(curr=curr, action=last_action))
         return last_action
-    def _run(self,shared): p=self.prep(shared); o=self._orch(shared); return self.post(shared,p,o)
-    def post(self,shared,prep_res,exec_res): return exec_res
+
+    def _run(self, shared: dict[str, Any]) -> Any:
+        """The execution wrapper for the flow.
+
+        This method wraps the `_orch` method with `prep` and `post` methods.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post` method.
+        """
+        p = self.prep(shared=shared)
+        o = self._orch(shared=shared)
+        return self.post(shared=shared, prep_res=p, exec_res=o)
+
+    def post(
+        self, shared: dict[str, Any], prep_res: Any, exec_res: Any
+    ) -> Any:
+        """The post-processing logic of the flow.
+
+        This method is called after the flow has finished executing.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+            prep_res (Any): The result of the `prep` method.
+            exec_res (Any): The result of the `_orch` method.
+
+        Returns:
+            Any: The result of the last node's `post` method.
+        """
+        return exec_res
+
 
 class BatchFlow(Flow):
-    def _run(self,shared):
-        pr=self.prep(shared) or []
-        for bp in pr: self._orch(shared,{**self.params,**bp})
-        return self.post(shared,pr,None)
+    """A `BatchFlow` object is a `Flow` that runs in batch mode.
+
+    It is a `Flow` that runs the same orchestration logic for a list of
+    different inputs using params
+
+    Example:
+
+    .. code-block:: python
+        
+        class MyBatchFlow(BatchFlow):
+            def prep(self, shared: dict[str, Any]) -> list[dict[str, Any]]:
+                # Prepare a list of inputs to process, possibly by looking at the shared storage
+                # must return a list of dicts
+                # each dict will be passed to the nodes in the flow as params
+                return [{'key1': 'data'}, {'key1': 'data2'}])
+            
+            # the nodes within the flow will operate on each item in the list
+            # by each looking at self.params
+            # and update the shared storage with the results
+
+            def post(self, shared: dict[str, Any], prep_res: list[dict[str, Any]], exec_res: None) -> Any:
+                # Examine shared storage to see the results of the nodes
+                return shared["results"]
+
+    """
+
+    def _run(self, shared: dict[str, Any]) -> Any:
+        """The execution wrapper for the batch flow.
+
+        This method wraps the `_orch` method with `prep` and `post` methods.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post` method.
+        """
+        pr = self.prep(shared=shared) or []
+        for bp in pr:
+            self._orch(shared=shared, params={**self.params, **bp})
+        return self.post(
+            shared=shared, prep_res=pr, exec_res=None
+        )
+
 
 class AsyncNode(Node):
-    async def prep_async(self,shared): pass
-    async def exec_async(self,prep_res): pass
-    async def exec_fallback_async(self,prep_res,exc): raise exc
-    async def post_async(self,shared,prep_res,exec_res): pass
-    async def _exec(self,prep_res): 
+    """An `AsyncNode` object is a `Node` that runs asynchronously.
+
+    It is a `Node` that can be used in an `AsyncFlow`.
+
+    Example:
+
+    .. code-block:: python
+        
+        class MyNode(AsyncNode):
+            def prep_async(self, shared: dict[str, Any]) -> Any:
+                # Prepare some data, often by looking at the shared storage
+                return {"data": shared.get("input", "default")}
+            def exec_async(self, prep_res: Any) -> Any:
+                # Execute some logic
+                return prep_res["data"].upper()
+            def post_async(self, shared: dict[str, Any], prep_res: Any, exec_res: Any) -> Any:
+                # Post-process the result
+                shared["result"] = exec_res
+                return "action_name"  # This will determine the next node in the flow
+                # returning None is the same as returning "default"
+    """
+
+    async def prep_async(self, shared: dict[str, Any]) -> Any:
+        """The async pre-processing logic of the node.
+
+        This method is called before the node is executed.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the pre-processing logic.
+        """
+        pass
+
+    async def exec_async(self, prep_res: Any) -> Any:
+        """The async execution logic of the node.
+
+        This method is called after the `prep_async` method.
+
+        Args:
+            prep_res (Any): The result of the `prep_async` method.
+
+        Returns:
+            Any: The result of the execution logic.
+        """
+        pass
+
+    async def exec_fallback_async(self, prep_res: Any, exc: Exception) -> Any:
+        """The async fallback mechanism for the node.
+
+        This method is called when the node has been retried `max_retries` times
+        and has failed every time.
+
+        Args:
+            prep_res (Any): The result of the `prep_async` method.
+            exc (Exception): The exception that was raised.
+
+        Raises:
+            exc: The exception that was raised.
+        """
+        raise exc
+
+    async def post_async(
+        self, shared: dict[str, Any], prep_res: Any, exec_res: Any
+    ) -> Any:
+        """The async post-processing logic of the node.
+
+        This method is called after the node is executed.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+            prep_res (Any): The result of the `prep_async` method.
+            exec_res (Any): The result of the `exec_async` method.
+
+        Returns:
+            Any: The result of the post-processing logic or an 'action'
+            returning None is the same as returning "default"
+        """
+        pass
+
+    async def _exec(self, prep_res: Any) -> Any:
+        """The async execution wrapper for the node.
+
+        This method wraps the `exec_async` method with retry and fallback logic.
+
+        Args:
+            prep_res (Any): The result of the `prep_async` method.
+
+        Returns:
+            Any: The result of the `exec_async` method or the
+                `exec_fallback_async` method.
+        """
         for i in range(self.max_retries):
-            try: return await self.exec_async(prep_res)
+            try:
+                return await self.exec_async(prep_res=prep_res)
             except Exception as e:
-                if i==self.max_retries-1: return await self.exec_fallback_async(prep_res,e)
-                if self.wait>0: await asyncio.sleep(self.wait)
-    async def run_async(self,shared): 
-        if self.successors: warnings.warn("Node won't run successors. Use AsyncFlow.")  
-        return await self._run_async(shared)
-    async def _run_async(self,shared): p=await self.prep_async(shared); e=await self._exec(p); return await self.post_async(shared,p,e)
-    def _run(self,shared): raise RuntimeError("Use run_async.")
-
-class AsyncBatchNode(AsyncNode,BatchNode):
-    async def _exec(self,items): return [await super(AsyncBatchNode,self)._exec(i) for i in items]
-
-class AsyncParallelBatchNode(AsyncNode,BatchNode):
-    async def _exec(self,items): return await asyncio.gather(*(super(AsyncParallelBatchNode,self)._exec(i) for i in items))
-
-class AsyncFlow(Flow,AsyncNode):
-    async def _orch_async(self,shared,params=None):
-        curr,p,last_action =copy.copy(self.start_node),(params or {**self.params}),None
-        while curr: curr.set_params(p); last_action=await curr._run_async(shared) if isinstance(curr,AsyncNode) else curr._run(shared); curr=copy.copy(self.get_next_node(curr,last_action))
+                if i == self.max_retries - 1:
+                    return await self.exec_fallback_async(
+                        prep_res=prep_res, exc=e
+                    )
+                if self.wait > 0:
+                    await asyncio.sleep(delay=self.wait)
+        return None  # Should be unreachable
+
+    async def run_async(self, shared: dict[str, Any]) -> Any:
+        """Runs the node asynchronously.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post_async` method.
+        """
+        if self.successors:
+            warnings.warn(message="Node won't run successors. Use AsyncFlow.")
+        return await self._run_async(shared=shared)
+
+    async def _run_async(self, shared: dict[str, Any]) -> Any:
+        """The async execution wrapper for the node.
+
+        This method wraps the `prep_async`, `_exec`, and `post_async` methods.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post_async` method.
+        """
+        p = await self.prep_async(shared=shared)
+        e = await self._exec(prep_res=p)
+        return await self.post_async(
+            shared=shared, prep_res=p, exec_res=e
+        )
+
+    def _run(self, shared: dict[str, Any]) -> Any:
+        """Raises a `RuntimeError` because this is an async node.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Raises:
+            RuntimeError: Always.
+        """
+        raise RuntimeError("Use run_async.")
+
+
+class AsyncBatchNode(AsyncNode, BatchNode):
+    """An `AsyncBatchNode` object is a `Node` that processes a batch of items
+    asynchronously.
+
+    The prep_async method should return a list of items to be processed
+    The exec_async method should process a single item and return a single result
+    The post_async method receives the list of results.
+    """
+
+    async def _exec(self, prep_res: list[Any] | None) -> list[Any]:  # type: ignore[override]
+        """The async execution wrapper for the batch node.
+
+        This method wraps the `_exec` method of the parent class to process a
+        batch of items.
+
+        Args:
+            prep_res (list[Any] | None): The result of the `prep` method.
+
+        Returns:
+            list[Any]: The list of results.
+        """
+        return [
+            await super(AsyncBatchNode, self)._exec(prep_res=i)
+            for i in prep_res or []
+        ]
+
+
+class AsyncParallelBatchNode(AsyncNode, BatchNode):
+    """An `AsyncParallelBatchNode` object is a `Node` that processes a batch of
+    items asynchronously and in parallel.
+
+    The prep_async method should return a list of items to be processed
+    The exec_async method should process each item in the list
+    The post_async method receivses the list of results.
+    """
+
+    async def _exec(self, prep_res: list[Any] | None) -> list[Any]:  # type: ignore[override]
+        """The async execution wrapper for the parallel batch node.
+
+        This method wraps the `_exec` method of the parent class to process a
+        batch of items in parallel.
+
+        Args:
+            prep_res (list[Any] | None): The result of the `prep` method.
+
+        Returns:
+            list[Any]: The list of results.
+        """
+        return await asyncio.gather(
+            *(
+                super(AsyncParallelBatchNode, self)._exec(prep_res=i)
+                for i in prep_res or []
+            )
+        )
+
+
+class AsyncFlow(Flow, AsyncNode):
+    """An `AsyncFlow` object is a `Flow` that runs asynchronously.
+
+    It is a `Flow` that can contain both `Node` and `AsyncNode` objects.
+
+    Example:
+
+    .. code-block:: python
+        
+        shared = {}
+        flow = AsyncFlow(start=AsyncNode())
+        ## the result of the last post_async method
+        result = await flow.run_async(shared=shared)
+        ## or examine an updated shared storage
+        print(shared)
+
+    """
+
+    async def _orch_async(
+        self, shared: dict[str, Any], params: dict[str, ParamValue] | None = None
+    ) -> str | None:
+        """The async orchestration logic of the flow.
+
+        This method executes the nodes in the flow in the order they are
+        connected.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+            params (dict[str, ParamValue] | None, optional): The parameters to pass to
+                the nodes. Defaults to None.
+
+        Returns:
+            str | None: The last action returned by a node in the flow.
+        """
+        curr, p, last_action = (
+            copy.copy(x=self.start_node),
+            (params or {**self.params}),
+            None,
+        )
+        while curr:
+            curr.set_params(params=p)
+            last_action = (
+                await curr._run_async(shared=shared)
+                if isinstance(curr, AsyncNode)
+                else curr._run(shared=shared)
+            )
+            curr = copy.copy(self.get_next_node(curr=curr, action=last_action))
         return last_action
-    async def _run_async(self,shared): p=await self.prep_async(shared); o=await self._orch_async(shared); return await self.post_async(shared,p,o)
-    async def post_async(self,shared,prep_res,exec_res): return exec_res
-
-class AsyncBatchFlow(AsyncFlow,BatchFlow):
-    async def _run_async(self,shared):
-        pr=await self.prep_async(shared) or []
-        for bp in pr: await self._orch_async(shared,{**self.params,**bp})
-        return await self.post_async(shared,pr,None)
-
-class AsyncParallelBatchFlow(AsyncFlow,BatchFlow):
-    async def _run_async(self,shared): 
-        pr=await self.prep_async(shared) or []
-        await asyncio.gather(*(self._orch_async(shared,{**self.params,**bp}) for bp in pr))
-        return await self.post_async(shared,pr,None)
\ No newline at end of file
+
+    async def _run_async(self, shared: dict[str, Any]) -> Any:
+        """The async execution wrapper for the flow.
+
+        This method wraps the `_orch_async` method with `prep_async` and
+        `post_async` methods.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post_async` method.
+        """
+        p = await self.prep_async(shared=shared)
+        o = await self._orch_async(shared=shared)
+        return await self.post_async(
+            shared=shared, prep_res=p, exec_res=o
+        )
+
+    async def post_async(
+        self, shared: dict[str, Any], prep_res: Any, exec_res: Any
+    ) -> Any:
+        """The async post-processing logic of the flow.
+
+        This method is called after the flow has finished executing.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+            prep_res (Any): The result of the `prep_async` method.
+            exec_res (Any): The result of the `_orch_async` method.
+
+        Returns:
+            Any: Often the result of the last node's `post_async` method
+        """
+        return exec_res
+
+
+class AsyncBatchFlow(AsyncFlow, BatchFlow):
+    """An `AsyncBatchFlow` object is a `Flow` that runs in batch mode
+    asynchronously.
+
+    It is a `Flow` that runs the same orchestration logic for a list of
+    different inputs.
+    """
+
+
+    async def _run_async(self, shared: dict[str, Any]) -> Any:
+        """The async execution wrapper for the batch flow.
+
+        This method wraps the `_orch_async` method with `prep_async` and
+        `post_async` methods.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post_async` method.
+        """
+        pr = await self.prep_async(shared=shared) or []
+        for bp in pr:
+            await self._orch_async(
+                shared=shared, params={**self.params, **bp}
+            )
+        return await self.post_async(
+            shared=shared, prep_res=pr, exec_res=None
+        )
+
+
+class AsyncParallelBatchFlow(AsyncFlow, BatchFlow):
+    """An `AsyncParallelBatchFlow` object is a `Flow` that runs in batch mode
+    asynchronously and in parallel.
+
+    It is a `Flow` that runs the same orchestration logic for a list of
+    different inputs.
+    """
+
+    async def _run_async(self, shared: dict[str, Any]) -> Any:
+        """The async execution wrapper for the parallel batch flow.
+
+        This method wraps the `_orch_async` method with `prep_async` and
+        `post_async` methods.
+
+        Args:
+            shared (dict[str, Any]): The shared storage of the flow.
+
+        Returns:
+            Any: The result of the `post_async` method.
+        """
+        pr = await self.prep_async(shared=shared) or []
+        await asyncio.gather(
+            *(
+                self._orch_async(
+                    shared=shared, params={**self.params, **bp}
+                )
+                for bp in pr
+            )
+        )
+        return await self.post_async(
+            shared=shared, prep_res=pr, exec_res=None
+        )
diff --git a/pocketflow/__init__.pyi b/pocketflow/__init__.pyi
deleted file mode 100644
index 220d2b4e..00000000
--- a/pocketflow/__init__.pyi
+++ /dev/null
@@ -1,97 +0,0 @@
-import asyncio
-from typing import Any, Dict, List, Optional, Union, TypeVar, Generic
-
-# Type variables for better type relationships
-_PrepResult = TypeVar('_PrepResult')
-_ExecResult = TypeVar('_ExecResult')
-_PostResult = TypeVar('_PostResult')
-
-# More specific parameter types
-ParamValue = Union[str, int, float, bool, None, List[Any], Dict[str, Any]]
-SharedData = Dict[str, Any]
-Params = Dict[str, ParamValue]
-
-class BaseNode(Generic[_PrepResult, _ExecResult, _PostResult]):
-    params: Params
-    successors: Dict[str, BaseNode[Any, Any, Any]]
-    
-    def __init__(self) -> None: ...
-    def set_params(self, params: Params) -> None: ...
-    def next(self, node: BaseNode[Any, Any, Any], action: str = "default") -> BaseNode[Any, Any, Any]: ...
-    def prep(self, shared: SharedData) -> _PrepResult: ...
-    def exec(self, prep_res: _PrepResult) -> _ExecResult: ...
-    def post(self, shared: SharedData, prep_res: _PrepResult, exec_res: _ExecResult) -> _PostResult: ...
-    def _exec(self, prep_res: _PrepResult) -> _ExecResult: ...
-    def _run(self, shared: SharedData) -> _PostResult: ...
-    def run(self, shared: SharedData) -> _PostResult: ...
-    def __rshift__(self, other: BaseNode[Any, Any, Any]) -> BaseNode[Any, Any, Any]: ...
-    def __sub__(self, action: str) -> _ConditionalTransition: ...
-
-class _ConditionalTransition:
-    src: BaseNode[Any, Any, Any]
-    action: str
-    
-    def __init__(self, src: BaseNode[Any, Any, Any], action: str) -> None: ...
-    def __rshift__(self, tgt: BaseNode[Any, Any, Any]) -> BaseNode[Any, Any, Any]: ...
-
-class Node(BaseNode[_PrepResult, _ExecResult, _PostResult]):
-    max_retries: int
-    wait: Union[int, float]
-    cur_retry: int
-    
-    def __init__(self, max_retries: int = 1, wait: Union[int, float] = 0) -> None: ...
-    def exec_fallback(self, prep_res: _PrepResult, exc: Exception) -> _ExecResult: ...
-    def _exec(self, prep_res: _PrepResult) -> _ExecResult: ...
-
-class BatchNode(Node[Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
-    def _exec(self, items: Optional[List[_PrepResult]]) -> List[_ExecResult]: ...
-
-class Flow(BaseNode[_PrepResult, Any, _PostResult]):
-    start_node: Optional[BaseNode[Any, Any, Any]]
-    
-    def __init__(self, start: Optional[BaseNode[Any, Any, Any]] = None) -> None: ...
-    def start(self, start: BaseNode[Any, Any, Any]) -> BaseNode[Any, Any, Any]: ...
-    def get_next_node(
-        self, curr: BaseNode[Any, Any, Any], action: Optional[str]
-    ) -> Optional[BaseNode[Any, Any, Any]]: ...
-    def _orch(
-        self, shared: SharedData, params: Optional[Params] = None
-    ) -> Any: ...
-    def _run(self, shared: SharedData) -> _PostResult: ...
-    def post(self, shared: SharedData, prep_res: _PrepResult, exec_res: Any) -> _PostResult: ...
-
-class BatchFlow(Flow[Optional[List[Params]], Any, _PostResult]):
-    def _run(self, shared: SharedData) -> _PostResult: ...
-
-class AsyncNode(Node[_PrepResult, _ExecResult, _PostResult]):
-    async def prep_async(self, shared: SharedData) -> _PrepResult: ...
-    async def exec_async(self, prep_res: _PrepResult) -> _ExecResult: ...
-    async def exec_fallback_async(self, prep_res: _PrepResult, exc: Exception) -> _ExecResult: ...
-    async def post_async(
-        self, shared: SharedData, prep_res: _PrepResult, exec_res: _ExecResult
-    ) -> _PostResult: ...
-    async def _exec(self, prep_res: _PrepResult) -> _ExecResult: ...
-    async def run_async(self, shared: SharedData) -> _PostResult: ...
-    async def _run_async(self, shared: SharedData) -> _PostResult: ...
-    def _run(self, shared: SharedData) -> _PostResult: ...
-
-class AsyncBatchNode(AsyncNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult], BatchNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
-    async def _exec(self, items: Optional[List[_PrepResult]]) -> List[_ExecResult]: ...
-
-class AsyncParallelBatchNode(AsyncNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult], BatchNode[Optional[List[_PrepResult]], List[_ExecResult], _PostResult]):
-    async def _exec(self, items: Optional[List[_PrepResult]]) -> List[_ExecResult]: ...
-
-class AsyncFlow(Flow[_PrepResult, Any, _PostResult], AsyncNode[_PrepResult, Any, _PostResult]):
-    async def _orch_async(
-        self, shared: SharedData, params: Optional[Params] = None
-    ) -> Any: ...
-    async def _run_async(self, shared: SharedData) -> _PostResult: ...
-    async def post_async(
-        self, shared: SharedData, prep_res: _PrepResult, exec_res: Any
-    ) -> _PostResult: ...
-
-class AsyncBatchFlow(AsyncFlow[Optional[List[Params]], Any, _PostResult], BatchFlow[Optional[List[Params]], Any, _PostResult]):
-    async def _run_async(self, shared: SharedData) -> _PostResult: ...
-
-class AsyncParallelBatchFlow(AsyncFlow[Optional[List[Params]], Any, _PostResult], BatchFlow[Optional[List[Params]], Any, _PostResult]):
-    async def _run_async(self, shared: SharedData) -> _PostResult: ...
\ No newline at end of file
diff --git a/tests/test_async_batch_flow.py b/tests/test_async_batch_flow.py
index 78f673b4..836b56ff 100644
--- a/tests/test_async_batch_flow.py
+++ b/tests/test_async_batch_flow.py
@@ -1,28 +1,29 @@
-import unittest
 import asyncio
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import AsyncNode, AsyncBatchFlow
+from pocketflow import AsyncBatchFlow, AsyncNode
+
 
 class AsyncDataProcessNode(AsyncNode):
-    async def prep_async(self, shared_storage):
+    async def prep_async(self, shared):
         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
+        data = shared['input_data'][key]
+        if 'results' not in shared:
+            shared['results'] = {}
+        shared['results'][key] = data
         return data
 
-    async def post_async(self, shared_storage, prep_result, proc_result):
+    async def post_async(self, shared, prep_res, exec_res):
         await asyncio.sleep(0.01)  # Simulate async work
         key = self.params.get('key')
-        shared_storage['results'][key] = prep_result * 2  # Double the value
+        shared['results'][key] = prep_res * 2  # Double the value
         return "processed"
 
 class AsyncErrorNode(AsyncNode):
-    async def post_async(self, shared_storage, prep_result, proc_result):
+    async def post_async(self, shared, prep_res, exec_res):
         key = self.params.get('key')
         if key == 'error_key':
             raise ValueError(f"Async error processing key: {key}")
@@ -35,10 +36,10 @@ def setUp(self):
     def test_basic_async_batch_processing(self):
         """Test basic async batch processing with multiple keys"""
         class SimpleTestAsyncBatchFlow(AsyncBatchFlow):
-            async def prep_async(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            async def prep_async(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'a': 1,
                 'b': 2,
@@ -47,37 +48,37 @@ async def prep_async(self, shared_storage):
         }
 
         flow = SimpleTestAsyncBatchFlow(start=self.process_node)
-        asyncio.run(flow.run_async(shared_storage))
+        asyncio.run(flow.run_async(shared))
 
         expected_results = {
             'a': 2,  # 1 * 2
             'b': 4,  # 2 * 2
             'c': 6   # 3 * 2
         }
-        self.assertEqual(shared_storage['results'], expected_results)
+        self.assertEqual(shared['results'], expected_results)
 
     def test_empty_async_batch(self):
         """Test async batch processing with empty input"""
         class EmptyTestAsyncBatchFlow(AsyncBatchFlow):
-            async def prep_async(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            async def prep_async(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
-        shared_storage = {
+        shared = {
             'input_data': {}
         }
 
         flow = EmptyTestAsyncBatchFlow(start=self.process_node)
-        asyncio.run(flow.run_async(shared_storage))
+        asyncio.run(flow.run_async(shared))
 
-        self.assertEqual(shared_storage.get('results', {}), {})
+        self.assertEqual(shared.get('results', {}), {})
 
     def test_async_error_handling(self):
         """Test error handling during async batch processing"""
         class ErrorTestAsyncBatchFlow(AsyncBatchFlow):
-            async def prep_async(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            async def prep_async(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'normal_key': 1,
                 'error_key': 2,
@@ -88,38 +89,38 @@ async def prep_async(self, shared_storage):
         flow = ErrorTestAsyncBatchFlow(start=AsyncErrorNode())
         
         with self.assertRaises(ValueError):
-            asyncio.run(flow.run_async(shared_storage))
+            asyncio.run(flow.run_async(shared))
 
     def test_nested_async_flow(self):
         """Test async batch processing with nested flows"""
         class AsyncInnerNode(AsyncNode):
-            async def post_async(self, shared_storage, prep_result, proc_result):
+            async def post_async(self, shared, prep_res, exec_res):
                 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
+                if 'intermediate_results' not in shared:
+                    shared['intermediate_results'] = {}
+                shared['intermediate_results'][key] = shared['input_data'][key] + 1
                 await asyncio.sleep(0.01)
                 return "next"
 
         class AsyncOuterNode(AsyncNode):
-            async def post_async(self, shared_storage, prep_result, proc_result):
+            async def post_async(self, shared, prep_res, exec_res):
                 key = self.params.get('key')
-                if 'results' not in shared_storage:
-                    shared_storage['results'] = {}
-                shared_storage['results'][key] = shared_storage['intermediate_results'][key] * 2
+                if 'results' not in shared:
+                    shared['results'] = {}
+                shared['results'][key] = shared['intermediate_results'][key] * 2
                 await asyncio.sleep(0.01)
                 return "done"
 
         class NestedAsyncBatchFlow(AsyncBatchFlow):
-            async def prep_async(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            async def prep_async(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
         # Create inner flow
         inner_node = AsyncInnerNode()
         outer_node = AsyncOuterNode()
         inner_node - "next" >> outer_node
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'x': 1,
                 'y': 2
@@ -127,34 +128,34 @@ async def prep_async(self, shared_storage):
         }
 
         flow = NestedAsyncBatchFlow(start=inner_node)
-        asyncio.run(flow.run_async(shared_storage))
+        asyncio.run(flow.run_async(shared))
 
         expected_results = {
             'x': 4,  # (1 + 1) * 2
             'y': 6   # (2 + 1) * 2
         }
-        self.assertEqual(shared_storage['results'], expected_results)
+        self.assertEqual(shared['results'], expected_results)
 
     def test_custom_async_parameters(self):
         """Test async batch processing with additional custom parameters"""
         class CustomParamAsyncNode(AsyncNode):
-            async def post_async(self, shared_storage, prep_result, proc_result):
+            async def post_async(self, shared, prep_res, exec_res):
                 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
+                if 'results' not in shared:
+                    shared['results'] = {}
+                shared['results'][key] = shared['input_data'][key] * multiplier
                 return "done"
 
         class CustomParamAsyncBatchFlow(AsyncBatchFlow):
-            async def prep_async(self, shared_storage):
+            async def prep_async(self, shared):
                 return [{
                     'key': k,
                     'multiplier': i + 1
-                } for i, k in enumerate(shared_storage['input_data'].keys())]
+                } for i, k in enumerate(shared['input_data'].keys())]
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'a': 1,
                 'b': 2,
@@ -163,14 +164,14 @@ async def prep_async(self, shared_storage):
         }
 
         flow = CustomParamAsyncBatchFlow(start=CustomParamAsyncNode())
-        asyncio.run(flow.run_async(shared_storage))
+        asyncio.run(flow.run_async(shared))
 
         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)
+        self.assertEqual(shared['results'], expected_results)
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/tests/test_async_batch_node.py b/tests/test_async_batch_node.py
index 2d3c8b01..21c76d1f 100644
--- a/tests/test_async_batch_node.py
+++ b/tests/test_async_batch_node.py
@@ -1,42 +1,43 @@
-import unittest
 import asyncio
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import AsyncNode, AsyncBatchNode, AsyncFlow
+from pocketflow import AsyncBatchNode, AsyncNode
+
 
 class AsyncArrayChunkNode(AsyncBatchNode):
     def __init__(self, chunk_size=10):
         super().__init__()
         self.chunk_size = chunk_size
     
-    async def prep_async(self, shared_storage):
+    async def prep_async(self, shared):
         # Get array from shared storage and split into chunks
-        array = shared_storage.get('input_array', [])
+        array = shared.get('input_array', [])
         chunks = []
         for start in range(0, len(array), self.chunk_size):
             end = min(start + self.chunk_size, len(array))
             chunks.append(array[start:end])
         return chunks
     
-    async def exec_async(self, chunk):
+    async def exec_async(self, prep_res):
         # Simulate async processing of each chunk
         await asyncio.sleep(0.01)
-        return sum(chunk)
+        return sum(prep_res)
         
-    async def post_async(self, shared_storage, prep_result, proc_result):
+    async def post_async(self, shared, prep_res, exec_res):
         # Store chunk results in shared storage
-        shared_storage['chunk_results'] = proc_result
+        shared['chunk_results'] = exec_res
         return "processed"
 
 class AsyncSumReduceNode(AsyncNode):
-    async def prep_async(self, shared_storage):
+    async def prep_async(self, shared):
         # Get chunk results from shared storage
-        chunk_results = shared_storage.get('chunk_results', [])
+        chunk_results = shared.get('chunk_results', [])
         await asyncio.sleep(0.01)  # Simulate async processing
         total = sum(chunk_results)
-        shared_storage['total'] = total
+        shared['total'] = total
         return "reduced"
 
 class TestAsyncBatchNode(unittest.TestCase):
@@ -44,14 +45,14 @@ def test_array_chunking(self):
         """
         Test that the array is correctly split into chunks and processed asynchronously
         """
-        shared_storage = {
+        shared = {
             'input_array': list(range(25))  # [0,1,2,...,24]
         }
         
         chunk_node = AsyncArrayChunkNode(chunk_size=10)
-        asyncio.run(chunk_node.run_async(shared_storage))
+        asyncio.run(chunk_node.run_async(shared=shared))
         
-        results = shared_storage['chunk_results']
+        results = shared['chunk_results']
         self.assertEqual(results, [45, 145, 110])  # Sum of chunks [0-9], [10-19], [20-24]
         
     # def test_async_map_reduce_sum(self):
@@ -63,7 +64,7 @@ def test_array_chunking(self):
     #     array = list(range(100))
     #     expected_sum = sum(array)  # 4950
         
-    #     shared_storage = {
+    #     shared = {
     #         'input_array': array
     #     }
         
@@ -76,9 +77,9 @@ def test_array_chunking(self):
         
     #     # Create and run pipeline
     #     pipeline = AsyncFlow(start=chunk_node)
-    #     asyncio.run(pipeline.run_async(shared_storage))
+    #     asyncio.run(pipeline.run_async(shared))
         
-    #     self.assertEqual(shared_storage['total'], expected_sum)
+    #     self.assertEqual(shared['total'], expected_sum)
         
     # def test_uneven_chunks(self):
     #     """
@@ -88,7 +89,7 @@ def test_array_chunking(self):
     #     array = list(range(25))
     #     expected_sum = sum(array)  # 300
         
-    #     shared_storage = {
+    #     shared = {
     #         'input_array': array
     #     }
         
@@ -97,9 +98,9 @@ def test_array_chunking(self):
         
     #     chunk_node - "processed" >> reduce_node
     #     pipeline = AsyncFlow(start=chunk_node)
-    #     asyncio.run(pipeline.run_async(shared_storage))
+    #     asyncio.run(pipeline.run_async(shared))
         
-    #     self.assertEqual(shared_storage['total'], expected_sum)
+    #     self.assertEqual(shared['total'], expected_sum)
 
     # def test_custom_chunk_size(self):
     #     """
@@ -108,7 +109,7 @@ def test_array_chunking(self):
     #     array = list(range(100))
     #     expected_sum = sum(array)
         
-    #     shared_storage = {
+    #     shared = {
     #         'input_array': array
     #     }
         
@@ -118,9 +119,9 @@ def test_array_chunking(self):
         
     #     chunk_node - "processed" >> reduce_node
     #     pipeline = AsyncFlow(start=chunk_node)
-    #     asyncio.run(pipeline.run_async(shared_storage))
+    #     asyncio.run(pipeline.run_async(shared))
         
-    #     self.assertEqual(shared_storage['total'], expected_sum)
+    #     self.assertEqual(shared['total'], expected_sum)
         
     # def test_single_element_chunks(self):
     #     """
@@ -129,7 +130,7 @@ def test_array_chunking(self):
     #     array = list(range(5))
     #     expected_sum = sum(array)
         
-    #     shared_storage = {
+    #     shared = {
     #         'input_array': array
     #     }
         
@@ -138,15 +139,15 @@ def test_array_chunking(self):
         
     #     chunk_node - "processed" >> reduce_node
     #     pipeline = AsyncFlow(start=chunk_node)
-    #     asyncio.run(pipeline.run_async(shared_storage))
+    #     asyncio.run(pipeline.run_async(shared))
         
-    #     self.assertEqual(shared_storage['total'], expected_sum)
+    #     self.assertEqual(shared['total'], expected_sum)
 
     # def test_empty_array(self):
     #     """
     #     Test edge case of empty input array
     #     """
-    #     shared_storage = {
+    #     shared = {
     #         'input_array': []
     #     }
         
@@ -155,27 +156,27 @@ def test_array_chunking(self):
         
     #     chunk_node - "processed" >> reduce_node
     #     pipeline = AsyncFlow(start=chunk_node)
-    #     asyncio.run(pipeline.run_async(shared_storage))
+    #     asyncio.run(pipeline.run_async(shared))
         
-    #     self.assertEqual(shared_storage['total'], 0)
+    #     self.assertEqual(shared['total'], 0)
 
     # def test_error_handling(self):
     #     """
     #     Test error handling in async batch processing
     #     """
     #     class ErrorAsyncBatchNode(AsyncBatchNode):
-    #         async def exec_async(self, item):
-    #             if item == 2:
+    #         async def exec_async(self, prep_res):
+    #             if prep_res == 2:
     #                 raise ValueError("Error processing item 2")
-    #             return item
+    #             return prep_res
 
-    #     shared_storage = {
+    #     shared = {
     #         'input_array': [1, 2, 3]
     #     }
         
     #     error_node = ErrorAsyncBatchNode()
     #     with self.assertRaises(ValueError):
-    #         asyncio.run(error_node.run_async(shared_storage))
+    #         asyncio.run(error_node.run_async(shared))
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/tests/test_async_flow.py b/tests/test_async_flow.py
index 38a9fbe3..00d9b400 100644
--- a/tests/test_async_flow.py
+++ b/tests/test_async_flow.py
@@ -1,10 +1,11 @@
-import unittest
 import asyncio
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import Node, AsyncNode, AsyncFlow
+from pocketflow import AsyncFlow, AsyncNode, Node
+
 
 class AsyncNumberNode(AsyncNode):
     """
@@ -16,13 +17,13 @@ def __init__(self, number):
         super().__init__()
         self.number = number
 
-    async def prep_async(self, shared_storage):
+    async def prep_async(self, shared):
         # Synchronous work is allowed inside an AsyncNode,
         # but final 'condition' is determined by post_async().
-        shared_storage['current'] = self.number
+        shared['current'] = self.number
         return "set_number"
 
-    async def post_async(self, shared_storage, prep_result, proc_result):
+    async def post_async(self, shared, prep_res, exec_res):
         # Possibly do asynchronous tasks here
         await asyncio.sleep(0.01)
         # Return a condition for the flow
@@ -32,11 +33,11 @@ class AsyncIncrementNode(AsyncNode):
     """
     Demonstrates incrementing the 'current' value asynchronously.
     """
-    async def prep_async(self, shared_storage):
-        shared_storage['current'] = shared_storage.get('current', 0) + 1
+    async def prep_async(self, shared):
+        shared['current'] = shared.get('current', 0) + 1
         return "incremented"
 
-    async def post_async(self, shared_storage, prep_result, proc_result):
+    async def post_async(self, shared, prep_res, exec_res):
         await asyncio.sleep(0.01)  # simulate async I/O
         return "done"
 
@@ -47,12 +48,12 @@ def __init__(self, signal="default_async_signal"):
         self.signal = signal
 
     # No prep needed usually if just signaling
-    async def prep_async(self, shared_storage):
+    async def prep_async(self, shared):
         await asyncio.sleep(0.01) # Simulate async work
 
-    async def post_async(self, shared_storage, prep_result, exec_result):
+    async def post_async(self, shared, prep_res, exec_res):
         # Store the signal in shared storage for verification
-        shared_storage['last_async_signal_emitted'] = self.signal
+        shared['last_async_signal_emitted'] = self.signal
         await asyncio.sleep(0.01) # Simulate async work
         print(self.signal)
         return self.signal # Return the specific action string
@@ -63,12 +64,12 @@ def __init__(self, path_id):
         super().__init__()
         self.path_id = path_id
 
-    async def prep_async(self, shared_storage):
+    async def prep_async(self, shared):
         await asyncio.sleep(0.01) # Simulate async work
-        shared_storage['async_path_taken'] = self.path_id
+        shared['async_path_taken'] = self.path_id
 
     # post_async implicitly returns None (for default transition out if needed)
-    async def post_async(self, shared_storage, prep_result, exec_result):
+    async def post_async(self, shared, prep_res, exec_res):
          await asyncio.sleep(0.01)
          # Return None by default
 
@@ -83,23 +84,23 @@ def test_async_number_node_direct_call(self):
         """
         async def run_node():
             node = AsyncNumberNode(42)
-            shared_storage = {}
-            condition = await node.run_async(shared_storage)
-            return shared_storage, condition
+            shared = {}
+            condition = await node.run_async(shared)
+            return shared, condition
 
-        shared_storage, condition = asyncio.run(run_node())
-        self.assertEqual(shared_storage['current'], 42)
+        shared, condition = asyncio.run(run_node())
+        self.assertEqual(shared['current'], 42)
         self.assertEqual(condition, "number_set")
 
     def test_async_increment_node_direct_call(self):
         async def run_node():
             node = AsyncIncrementNode()
-            shared_storage = {'current': 10}
-            condition = await node.run_async(shared_storage)
-            return shared_storage, condition
+            shared = {'current': 10}
+            condition = await node.run_async(shared)
+            return shared, condition
 
-        shared_storage, condition = asyncio.run(run_node())
-        self.assertEqual(shared_storage['current'], 11)
+        shared, condition = asyncio.run(run_node())
+        self.assertEqual(shared['current'], 11)
         self.assertEqual(condition, "done")
 
 
@@ -125,10 +126,10 @@ def test_simple_async_flow(self):
         flow = AsyncFlow(start)
 
         # We'll run the flow synchronously (which under the hood is asyncio.run())
-        shared_storage = {}
-        asyncio.run(flow.run_async(shared_storage))
+        shared = {}
+        asyncio.run(flow.run_async(shared))
 
-        self.assertEqual(shared_storage['current'], 6)
+        self.assertEqual(shared['current'], 6)
 
     def test_async_flow_branching(self):
         """
@@ -139,30 +140,30 @@ def test_async_flow_branching(self):
         """
 
         class BranchingAsyncNode(AsyncNode):
-            def exec(self, data):
-                value = shared_storage.get("value", 0)
-                shared_storage["value"] = value
+            def exec(self, prep_res):
+                value = shared.get("value", 0)
+                shared["value"] = value
                 # We'll decide branch based on whether 'value' is positive
                 return None
 
-            async def post_async(self, shared_storage, prep_result, proc_result):
+            async def post_async(self, shared, prep_res, exec_res):
                 await asyncio.sleep(0.01)
-                if shared_storage["value"] >= 0:
+                if shared["value"] >= 0:
                     return "positive_branch"
                 else:
                     return "negative_branch"
 
         class PositiveNode(Node):
-            def exec(self, data):
-                shared_storage["path"] = "positive"
+            def exec(self, prep_res):
+                shared["path"] = "positive"
                 return None
 
         class NegativeNode(Node):
-            def exec(self, data):
-                shared_storage["path"] = "negative"
+            def exec(self, prep_res):
+                shared["path"] = "negative"
                 return None
 
-        shared_storage = {"value": 10}
+        shared = {"value": 10}
 
         start = BranchingAsyncNode()
         positive_node = PositiveNode()
@@ -173,9 +174,9 @@ def exec(self, data):
         start - "negative_branch" >> negative_node
 
         flow = AsyncFlow(start)
-        asyncio.run(flow.run_async(shared_storage))
+        asyncio.run(flow.run_async(shared))
 
-        self.assertEqual(shared_storage["path"], "positive", 
+        self.assertEqual(shared["path"], "positive", 
                          "Should have taken the positive branch")
 
     def test_async_composition_with_action_propagation(self):
@@ -183,7 +184,7 @@ def test_async_composition_with_action_propagation(self):
         Test AsyncFlow branches based on action from nested AsyncFlow's last node.
         """
         async def run_test():
-            shared_storage = {}
+            shared = {}
 
             # 1. Define an inner async flow ending with AsyncSignalNode
             # Use existing AsyncNumberNode which should return None from post_async implicitly
@@ -206,20 +207,20 @@ async def run_test():
 
             # 5. Run the outer async flow and capture the last action
             # Execution: inner_start -> inner_end -> path_b
-            last_action_outer = await outer_flow.run_async(shared_storage)
+            last_action_outer = await outer_flow.run_async(shared)
 
             # 6. Return results for assertion
-            return shared_storage, last_action_outer
+            return shared, last_action_outer
 
         # Run the async test function
-        shared_storage, last_action_outer = asyncio.run(run_test())
+        shared, last_action_outer = asyncio.run(run_test())
 
         # 7. Assert the results
         # Check state after inner flow execution
-        self.assertEqual(shared_storage.get('current'), 200) # From AsyncNumberNode
-        self.assertEqual(shared_storage.get('last_async_signal_emitted'), "async_inner_done")
+        self.assertEqual(shared.get('current'), 200) # From AsyncNumberNode
+        self.assertEqual(shared.get('last_async_signal_emitted'), "async_inner_done")
         # Check that the correct outer path was taken
-        self.assertEqual(shared_storage.get('async_path_taken'), "AsyncB")
+        self.assertEqual(shared.get('async_path_taken'), "AsyncB")
         # Check the action returned by the outer flow. The last node executed was
         # path_b_node, which returns None from its post_async method.
         self.assertIsNone(last_action_outer)
diff --git a/tests/test_async_parallel_batch_flow.py b/tests/test_async_parallel_batch_flow.py
index ab522fc3..bdccaedb 100644
--- a/tests/test_async_parallel_batch_flow.py
+++ b/tests/test_async_parallel_batch_flow.py
@@ -1,45 +1,46 @@
-import unittest
 import asyncio
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import AsyncNode, AsyncParallelBatchNode, AsyncParallelBatchFlow
+from pocketflow import AsyncNode, AsyncParallelBatchFlow, AsyncParallelBatchNode
+
 
 class AsyncParallelNumberProcessor(AsyncParallelBatchNode):
     def __init__(self, delay=0.1):
         super().__init__()
         self.delay = delay
     
-    async def prep_async(self, shared_storage):
-        batch = shared_storage['batches'][self.params['batch_id']]
+    async def prep_async(self, shared):
+        batch = shared['batches'][self.params['batch_id']]
         return batch
     
-    async def exec_async(self, number):
+    async def exec_async(self, prep_res):
         await asyncio.sleep(self.delay)  # Simulate async processing
-        return number * 2
+        return prep_res * 2
         
-    async def post_async(self, shared_storage, prep_result, exec_result):
-        if 'processed_numbers' not in shared_storage:
-            shared_storage['processed_numbers'] = {}
-        shared_storage['processed_numbers'][self.params['batch_id']] = exec_result
+    async def post_async(self, shared, prep_res, exec_res):
+        if 'processed_numbers' not in shared:
+            shared['processed_numbers'] = {}
+        shared['processed_numbers'][self.params['batch_id']] = exec_res
         return "processed"
 
 class AsyncAggregatorNode(AsyncNode):
-    async def prep_async(self, shared_storage):
-        # Combine all batch results in order
+    async def prep_async(self, shared):
         all_results = []
-        processed = shared_storage.get('processed_numbers', {})
-        for i in range(len(processed)):
+        processed = shared.get('processed_numbers', {})
+        # To maintain the original order, sort by batch_id (the keys)
+        for i in sorted(processed.keys()):
             all_results.extend(processed[i])
         return all_results
     
-    async def exec_async(self, prep_result):
+    async def exec_async(self, prep_res):
         await asyncio.sleep(0.01)
-        return sum(prep_result)
+        return sum(prep_res)
     
-    async def post_async(self, shared_storage, prep_result, exec_result):
-        shared_storage['total'] = exec_result
+    async def post_async(self, shared, prep_res, exec_res):
+        shared['total'] = exec_res
         return "aggregated"
 
 class TestAsyncParallelBatchFlow(unittest.TestCase):
@@ -55,10 +56,10 @@ def test_parallel_batch_flow(self):
         Test basic parallel batch processing flow with batch IDs
         """
         class TestParallelBatchFlow(AsyncParallelBatchFlow):
-            async def prep_async(self, shared_storage):
-                return [{'batch_id': i} for i in range(len(shared_storage['batches']))]
+            async def prep_async(self, shared):
+                return [{'batch_id': i} for i in range(len(shared['batches']))]
 
-        shared_storage = {
+        shared = {
             'batches': [
                 [1, 2, 3],  # batch_id: 0
                 [4, 5, 6],  # batch_id: 1
@@ -73,7 +74,7 @@ async def prep_async(self, shared_storage):
         flow = TestParallelBatchFlow(start=processor)
         
         start_time = self.loop.time()
-        self.loop.run_until_complete(flow.run_async(shared_storage))
+        self.loop.run_until_complete(flow.run_async(shared))
         execution_time = self.loop.time() - start_time
 
         # Verify each batch was processed correctly
@@ -82,11 +83,11 @@ async def prep_async(self, shared_storage):
             1: [8, 10, 12],  # [4,5,6] * 2
             2: [14, 16, 18]  # [7,8,9] * 2
         }
-        self.assertEqual(shared_storage['processed_numbers'], expected_batch_results)
+        self.assertEqual(shared['processed_numbers'], expected_batch_results)
         
         # Verify total
-        expected_total = sum(num * 2 for batch in shared_storage['batches'] for num in batch)
-        self.assertEqual(shared_storage['total'], expected_total)
+        expected_total = sum(num * 2 for batch in shared['batches'] for num in batch)
+        self.assertEqual(shared['total'], expected_total)
         
         # Verify parallel execution
         self.assertLess(execution_time, 0.2)
@@ -96,16 +97,16 @@ def test_error_handling(self):
         Test error handling in parallel batch flow
         """
         class ErrorProcessor(AsyncParallelNumberProcessor):
-            async def exec_async(self, item):
-                if item == 2:
-                    raise ValueError(f"Error processing item {item}")
-                return item
+            async def exec_async(self, prep_res):
+                if prep_res == 2:
+                    raise ValueError(f"Error processing item {prep_res}")
+                return prep_res
 
         class ErrorBatchFlow(AsyncParallelBatchFlow):
-            async def prep_async(self, shared_storage):
-                return [{'batch_id': i} for i in range(len(shared_storage['batches']))]
+            async def prep_async(self, shared):
+                return [{'batch_id': i} for i in range(len(shared['batches']))]
 
-        shared_storage = {
+        shared = {
             'batches': [
                 [1, 2, 3],  # Contains error-triggering value
                 [4, 5, 6]
@@ -116,17 +117,17 @@ async def prep_async(self, shared_storage):
         flow = ErrorBatchFlow(start=processor)
         
         with self.assertRaises(ValueError):
-            self.loop.run_until_complete(flow.run_async(shared_storage))
+            self.loop.run_until_complete(flow.run_async(shared))
 
     def test_multiple_batch_sizes(self):
         """
         Test parallel batch flow with varying batch sizes
         """
         class VaryingBatchFlow(AsyncParallelBatchFlow):
-            async def prep_async(self, shared_storage):
-                return [{'batch_id': i} for i in range(len(shared_storage['batches']))]
+            async def prep_async(self, shared):
+                return [{'batch_id': i} for i in range(len(shared['batches']))]
 
-        shared_storage = {
+        shared = {
             'batches': [
                 [1],           # batch_id: 0
                 [2, 3, 4],    # batch_id: 1
@@ -141,7 +142,7 @@ async def prep_async(self, shared_storage):
         processor - "processed" >> aggregator
         flow = VaryingBatchFlow(start=processor)
         
-        self.loop.run_until_complete(flow.run_async(shared_storage))
+        self.loop.run_until_complete(flow.run_async(shared))
         
         # Verify each batch was processed correctly
         expected_batch_results = {
@@ -150,11 +151,11 @@ async def prep_async(self, shared_storage):
             2: [10, 12],           # [5,6] * 2
             3: [14, 16, 18, 20]    # [7,8,9,10] * 2
         }
-        self.assertEqual(shared_storage['processed_numbers'], expected_batch_results)
+        self.assertEqual(shared['processed_numbers'], expected_batch_results)
         
         # Verify total
-        expected_total = sum(num * 2 for batch in shared_storage['batches'] for num in batch)
-        self.assertEqual(shared_storage['total'], expected_total)
+        expected_total = sum(num * 2 for batch in shared['batches'] for num in batch)
+        self.assertEqual(shared['total'], expected_total)
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/tests/test_async_parallel_batch_node.py b/tests/test_async_parallel_batch_node.py
index 34d9aced..a699a4bc 100644
--- a/tests/test_async_parallel_batch_node.py
+++ b/tests/test_async_parallel_batch_node.py
@@ -1,26 +1,27 @@
-import unittest
 import asyncio
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import AsyncParallelBatchNode, AsyncParallelBatchFlow
+from pocketflow import AsyncParallelBatchNode
+
 
 class AsyncParallelNumberProcessor(AsyncParallelBatchNode):
     def __init__(self, delay=0.1):
         super().__init__()
         self.delay = delay
     
-    async def prep_async(self, shared_storage):
-        numbers = shared_storage.get('input_numbers', [])
+    async def prep_async(self, shared):
+        numbers = shared.get('input_numbers', [])
         return numbers
     
-    async def exec_async(self, number):
+    async def exec_async(self, prep_res):
         await asyncio.sleep(self.delay)  # Simulate async processing
-        return number * 2
+        return prep_res * 2
         
-    async def post_async(self, shared_storage, prep_result, exec_result):
-        shared_storage['processed_numbers'] = exec_result
+    async def post_async(self, shared, prep_res, exec_res):
+        shared['processed_numbers'] = exec_res
         return "processed"
 
 class TestAsyncParallelBatchNode(unittest.TestCase):
@@ -36,7 +37,7 @@ def test_parallel_processing(self):
         """
         Test that numbers are processed in parallel by measuring execution time
         """
-        shared_storage = {
+        shared = {
             'input_numbers': list(range(5))
         }
         
@@ -44,12 +45,12 @@ def test_parallel_processing(self):
         
         # Run the processor
         start_time = asyncio.get_event_loop().time()
-        self.loop.run_until_complete(processor.run_async(shared_storage))
+        self.loop.run_until_complete(processor.run_async(shared))
         end_time = asyncio.get_event_loop().time()
         
         # Check results
         expected = [0, 2, 4, 6, 8]  # Each number doubled
-        self.assertEqual(shared_storage['processed_numbers'], expected)
+        self.assertEqual(shared['processed_numbers'], expected)
         
         # Since processing is parallel, total time should be approximately
         # equal to the delay of a single operation, not delay * number_of_items
@@ -60,60 +61,60 @@ def test_empty_input(self):
         """
         Test processing of empty input
         """
-        shared_storage = {
+        shared = {
             'input_numbers': []
         }
         
         processor = AsyncParallelNumberProcessor()
-        self.loop.run_until_complete(processor.run_async(shared_storage))
+        self.loop.run_until_complete(processor.run_async(shared))
         
-        self.assertEqual(shared_storage['processed_numbers'], [])
+        self.assertEqual(shared['processed_numbers'], [])
     
     def test_single_item(self):
         """
         Test processing of a single item
         """
-        shared_storage = {
+        shared = {
             'input_numbers': [42]
         }
         
         processor = AsyncParallelNumberProcessor()
-        self.loop.run_until_complete(processor.run_async(shared_storage))
+        self.loop.run_until_complete(processor.run_async(shared))
         
-        self.assertEqual(shared_storage['processed_numbers'], [84])
+        self.assertEqual(shared['processed_numbers'], [84])
     
     def test_large_batch(self):
         """
         Test processing of a large batch of numbers
         """
         input_size = 100
-        shared_storage = {
+        shared = {
             'input_numbers': list(range(input_size))
         }
         
         processor = AsyncParallelNumberProcessor(delay=0.01)
-        self.loop.run_until_complete(processor.run_async(shared_storage))
+        self.loop.run_until_complete(processor.run_async(shared))
         
         expected = [x * 2 for x in range(input_size)]
-        self.assertEqual(shared_storage['processed_numbers'], expected)
+        self.assertEqual(shared['processed_numbers'], expected)
     
     def test_error_handling(self):
         """
         Test error handling during parallel processing
         """
         class ErrorProcessor(AsyncParallelNumberProcessor):
-            async def exec_async(self, item):
-                if item == 2:
-                    raise ValueError(f"Error processing item {item}")
-                return item
+            async def exec_async(self, prep_res):
+                if prep_res == 2:
+                    raise ValueError(f"Error processing item {prep_res}")
+                return prep_res
         
-        shared_storage = {
+        shared = {
             'input_numbers': [1, 2, 3]
         }
         
         processor = ErrorProcessor()
         with self.assertRaises(ValueError):
-            self.loop.run_until_complete(processor.run_async(shared_storage))
+            self.loop.run_until_complete(processor.run_async(shared))
     
     def test_concurrent_execution(self):
         """
@@ -122,18 +123,18 @@ def test_concurrent_execution(self):
         execution_order = []
         
         class OrderTrackingProcessor(AsyncParallelNumberProcessor):
-            async def exec_async(self, item):
-                delay = 0.1 if item % 2 == 0 else 0.05
+            async def exec_async(self, prep_res):
+                delay = 0.1 if prep_res % 2 == 0 else 0.05
                 await asyncio.sleep(delay)
-                execution_order.append(item)
-                return item
+                execution_order.append(prep_res)
+                return prep_res
         
-        shared_storage = {
+        shared = {
             'input_numbers': list(range(4))  # [0, 1, 2, 3]
         }
         
         processor = OrderTrackingProcessor()
-        self.loop.run_until_complete(processor.run_async(shared_storage))
+        self.loop.run_until_complete(processor.run_async(shared))
         
         # Odd numbers should finish before even numbers due to shorter delay
         self.assertLess(execution_order.index(1), execution_order.index(0))
diff --git a/tests/test_batch_flow.py b/tests/test_batch_flow.py
index d5ae426b..95201e1a 100644
--- a/tests/test_batch_flow.py
+++ b/tests/test_batch_flow.py
@@ -1,26 +1,27 @@
-import unittest
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import Node, BatchFlow, Flow
+from pocketflow import BatchFlow, Node
+
 
 class DataProcessNode(Node):
-    def prep(self, shared_storage):
+    def prep(self, shared):
         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 * 2
+        data = shared['input_data'][key]
+        if 'results' not in shared:
+            shared['results'] = {}
+        shared['results'][key] = data * 2
 
 class ErrorProcessNode(Node):
-    def prep(self, shared_storage):
+    def prep(self, shared):
         key = self.params.get('key')
         if key == 'error_key':
             raise ValueError(f"Error processing key: {key}")
-        if 'results' not in shared_storage:
-            shared_storage['results'] = {}
-        shared_storage['results'][key] = True
+        if 'results' not in shared:
+            shared['results'] = {}
+        shared['results'][key] = True
 
 class TestBatchFlow(unittest.TestCase):
     def setUp(self):
@@ -29,10 +30,10 @@ def setUp(self):
     def test_basic_batch_processing(self):
         """Test basic batch processing with multiple keys"""
         class SimpleTestBatchFlow(BatchFlow):
-            def prep(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            def prep(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'a': 1,
                 'b': 2,
@@ -41,57 +42,57 @@ def prep(self, shared_storage):
         }
 
         flow = SimpleTestBatchFlow(start=self.process_node)
-        flow.run(shared_storage)
+        flow.run(shared)
 
         expected_results = {
             'a': 2,
             'b': 4,
             'c': 6
         }
-        self.assertEqual(shared_storage['results'], expected_results)
+        self.assertEqual(shared['results'], expected_results)
 
     def test_empty_input(self):
         """Test batch processing with empty input dictionary"""
         class EmptyTestBatchFlow(BatchFlow):
-            def prep(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            def prep(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
-        shared_storage = {
+        shared = {
             'input_data': {}
         }
 
         flow = EmptyTestBatchFlow(start=self.process_node)
-        flow.run(shared_storage)
+        flow.run(shared)
 
-        self.assertEqual(shared_storage.get('results', {}), {})
+        self.assertEqual(shared.get('results', {}), {})
 
     def test_single_item(self):
         """Test batch processing with single item"""
         class SingleItemBatchFlow(BatchFlow):
-            def prep(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            def prep(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'single': 5
             }
         }
 
         flow = SingleItemBatchFlow(start=self.process_node)
-        flow.run(shared_storage)
+        flow.run(shared)
 
         expected_results = {
             'single': 10
         }
-        self.assertEqual(shared_storage['results'], expected_results)
+        self.assertEqual(shared['results'], expected_results)
 
     def test_error_handling(self):
         """Test error handling during batch processing"""
         class ErrorTestBatchFlow(BatchFlow):
-            def prep(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            def prep(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'normal_key': 1,
                 'error_key': 2,
@@ -102,34 +103,34 @@ def prep(self, shared_storage):
         flow = ErrorTestBatchFlow(start=ErrorProcessNode())
         
         with self.assertRaises(ValueError):
-            flow.run(shared_storage)
+            flow.run(shared)
 
     def test_nested_flow(self):
         """Test batch processing with nested flows"""
         class InnerNode(Node):
-            def exec(self, prep_result):
+            def exec(self, prep_res):
                 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
+                if 'intermediate_results' not in shared:
+                    shared['intermediate_results'] = {}
+                shared['intermediate_results'][key] = shared['input_data'][key] + 1
 
         class OuterNode(Node):
-            def exec(self, prep_result):
+            def exec(self, prep_res):
                 key = self.params.get('key')
-                if 'results' not in shared_storage:
-                    shared_storage['results'] = {}
-                shared_storage['results'][key] = shared_storage['intermediate_results'][key] * 2
+                if 'results' not in shared:
+                    shared['results'] = {}
+                shared['results'][key] = shared['intermediate_results'][key] * 2
 
         class NestedBatchFlow(BatchFlow):
-            def prep(self, shared_storage):
-                return [{'key': k} for k in shared_storage['input_data'].keys()]
+            def prep(self, shared):
+                return [{'key': k} for k in shared['input_data'].keys()]
 
         # Create inner flow
         inner_node = InnerNode()
         outer_node = OuterNode()
         inner_node >> outer_node
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'x': 1,
                 'y': 2
@@ -137,32 +138,32 @@ def prep(self, shared_storage):
         }
 
         flow = NestedBatchFlow(start=inner_node)
-        flow.run(shared_storage)
+        flow.run(shared)
 
         expected_results = {
             'x': 4,  # (1 + 1) * 2
             'y': 6   # (2 + 1) * 2
         }
-        self.assertEqual(shared_storage['results'], expected_results)
+        self.assertEqual(shared['results'], expected_results)
 
     def test_custom_parameters(self):
         """Test batch processing with additional custom parameters"""
         class CustomParamNode(Node):
-            def exec(self, prep_result):
+            def exec(self, prep_res):
                 key = self.params.get('key')
                 multiplier = self.params.get('multiplier', 1)
-                if 'results' not in shared_storage:
-                    shared_storage['results'] = {}
-                shared_storage['results'][key] = shared_storage['input_data'][key] * multiplier
+                if 'results' not in shared:
+                    shared['results'] = {}
+                shared['results'][key] = shared['input_data'][key] * multiplier
 
         class CustomParamBatchFlow(BatchFlow):
-            def prep(self, shared_storage):
+            def prep(self, shared):
                 return [{
                     'key': k,
                     'multiplier': i + 1
-                } for i, k in enumerate(shared_storage['input_data'].keys())]
+                } for i, k in enumerate(shared['input_data'].keys())]
 
-        shared_storage = {
+        shared = {
             'input_data': {
                 'a': 1,
                 'b': 2,
@@ -171,14 +172,14 @@ def prep(self, shared_storage):
         }
 
         flow = CustomParamBatchFlow(start=CustomParamNode())
-        flow.run(shared_storage)
+        flow.run(shared)
 
         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)
+        self.assertEqual(shared['results'], expected_results)
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/tests/test_batch_node.py b/tests/test_batch_node.py
index 7e10dc52..70dba02a 100644
--- a/tests/test_batch_node.py
+++ b/tests/test_batch_node.py
@@ -1,53 +1,54 @@
-import unittest
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import Node, BatchNode, Flow
+from pocketflow import BatchNode, Flow, Node
+
 
 class ArrayChunkNode(BatchNode):
     def __init__(self, chunk_size=10):
         super().__init__()
         self.chunk_size = chunk_size
     
-    def prep(self, shared_storage):
+    def prep(self, shared):
         # Get array from shared storage and split into chunks
-        array = shared_storage.get('input_array', [])
+        array = shared.get('input_array', [])
         chunks = []
         for start in range(0, len(array), self.chunk_size):
             end = min(start + self.chunk_size, len(array))
             chunks.append(array[start: end])
         return chunks
     
-    def exec(self, chunk):
+    def exec(self, prep_res):
         # Process the chunk and return its sum
-        chunk_sum = sum(chunk)
+        chunk_sum = sum(prep_res)
         return chunk_sum
         
-    def post(self, shared_storage, prep_result, proc_result):
+    def post(self, shared, prep_res, exec_res):
         # Store chunk results in shared storage
-        shared_storage['chunk_results'] = proc_result
+        shared['chunk_results'] = exec_res
         return "default"
 
 class SumReduceNode(Node):
-    def prep(self, shared_storage):
+    def prep(self, shared):
         # Get chunk results from shared storage and sum them
-        chunk_results = shared_storage.get('chunk_results', [])
+        chunk_results = shared.get('chunk_results', [])
         total = sum(chunk_results)
-        shared_storage['total'] = total
+        shared['total'] = total
 
 class TestBatchNode(unittest.TestCase):
     def test_array_chunking(self):
         """
         Test that the array is correctly split into chunks
         """
-        shared_storage = {
+        shared = {
             'input_array': list(range(25))  # [0,1,2,...,24]
         }
         
         chunk_node = ArrayChunkNode(chunk_size=10)
-        chunk_node.run(shared_storage)
-        results = shared_storage['chunk_results']
+        chunk_node.run(shared)
+        results = shared['chunk_results']
         self.assertEqual(results, [45, 145, 110])
         
     def test_map_reduce_sum(self):
@@ -60,7 +61,7 @@ def test_map_reduce_sum(self):
         array = list(range(100))
         expected_sum = sum(array)  # 4950
         
-        shared_storage = {
+        shared = {
             'input_array': array
         }
         
@@ -73,9 +74,9 @@ def test_map_reduce_sum(self):
         
         # Create and run pipeline
         pipeline = Flow(start=chunk_node)
-        pipeline.run(shared_storage)
+        pipeline.run(shared)
         
-        self.assertEqual(shared_storage['total'], expected_sum)
+        self.assertEqual(shared['total'], expected_sum)
         
     def test_uneven_chunks(self):
         """
@@ -85,7 +86,7 @@ def test_uneven_chunks(self):
         array = list(range(25))
         expected_sum = sum(array)  # 300
         
-        shared_storage = {
+        shared = {
             'input_array': array
         }
         
@@ -94,9 +95,9 @@ def test_uneven_chunks(self):
         
         chunk_node >> reduce_node
         pipeline = Flow(start=chunk_node)
-        pipeline.run(shared_storage)
+        pipeline.run(shared)
         
-        self.assertEqual(shared_storage['total'], expected_sum)
+        self.assertEqual(shared['total'], expected_sum)
 
     def test_custom_chunk_size(self):
         """
@@ -105,7 +106,7 @@ def test_custom_chunk_size(self):
         array = list(range(100))
         expected_sum = sum(array)
         
-        shared_storage = {
+        shared = {
             'input_array': array
         }
         
@@ -115,9 +116,9 @@ def test_custom_chunk_size(self):
         
         chunk_node >> reduce_node
         pipeline = Flow(start=chunk_node)
-        pipeline.run(shared_storage)
+        pipeline.run(shared)
         
-        self.assertEqual(shared_storage['total'], expected_sum)
+        self.assertEqual(shared['total'], expected_sum)
         
     def test_single_element_chunks(self):
         """
@@ -126,7 +127,7 @@ def test_single_element_chunks(self):
         array = list(range(5))
         expected_sum = sum(array)
         
-        shared_storage = {
+        shared = {
             'input_array': array
         }
         
@@ -135,15 +136,15 @@ def test_single_element_chunks(self):
         
         chunk_node >> reduce_node
         pipeline = Flow(start=chunk_node)
-        pipeline.run(shared_storage)
+        pipeline.run(shared)
         
-        self.assertEqual(shared_storage['total'], expected_sum)
+        self.assertEqual(shared['total'], expected_sum)
 
     def test_empty_array(self):
         """
         Test edge case of empty input array
         """
-        shared_storage = {
+        shared = {
             'input_array': []
         }
         
@@ -152,9 +153,9 @@ def test_empty_array(self):
         
         chunk_node >> reduce_node
         pipeline = Flow(start=chunk_node)
-        pipeline.run(shared_storage)
+        pipeline.run(shared)
         
-        self.assertEqual(shared_storage['total'], 0)
+        self.assertEqual(shared['total'], 0)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_fall_back.py b/tests/test_fall_back.py
index 6d455e68..2d1aed41 100644
--- a/tests/test_fall_back.py
+++ b/tests/test_fall_back.py
@@ -1,10 +1,11 @@
-import unittest
 import asyncio
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import Node, AsyncNode, Flow, AsyncFlow
+from pocketflow import AsyncFlow, AsyncNode, Flow, Node
+
 
 class FallbackNode(Node):
     def __init__(self, should_fail=True, max_retries=1):
@@ -12,24 +13,24 @@ def __init__(self, should_fail=True, max_retries=1):
         self.should_fail = should_fail
         self.attempt_count = 0
     
-    def prep(self, shared_storage):
-        if 'results' not in shared_storage:
-            shared_storage['results'] = []
+    def prep(self, shared):
+        if 'results' not in shared:
+            shared['results'] = []
         return None
     
-    def exec(self, prep_result):
+    def exec(self, prep_res):
         self.attempt_count += 1
         if self.should_fail:
             raise ValueError("Intentional failure")
         return "success"
     
-    def exec_fallback(self, prep_result, exc):
+    def exec_fallback(self, prep_res, exc):
         return "fallback"
     
-    def post(self, shared_storage, prep_result, exec_result):
-        shared_storage['results'].append({
+    def post(self, shared, prep_res, exec_res):
+        shared['results'].append({
             'attempts': self.attempt_count,
-            'result': exec_result
+            'result': exec_res
         })
 
 class AsyncFallbackNode(AsyncNode):
@@ -38,102 +39,102 @@ def __init__(self, should_fail=True, max_retries=1):
         self.should_fail = should_fail
         self.attempt_count = 0
     
-    async def prep_async(self, shared_storage):
-        if 'results' not in shared_storage:
-            shared_storage['results'] = []
+    async def prep_async(self, shared):
+        if 'results' not in shared:
+            shared['results'] = []
         return None
     
-    async def exec_async(self, prep_result):
+    async def exec_async(self, prep_res):
         self.attempt_count += 1
         if self.should_fail:
             raise ValueError("Intentional async failure")
         return "success"
     
-    async def exec_fallback_async(self, prep_result, exc):
+    async def exec_fallback_async(self, prep_res, exc):
         await asyncio.sleep(0.01)  # Simulate async work
         return "async_fallback"
     
-    async def post_async(self, shared_storage, prep_result, exec_result):
-        shared_storage['results'].append({
+    async def post_async(self, shared, prep_res, exec_res):
+        shared['results'].append({
             'attempts': self.attempt_count,
-            'result': exec_result
+            'result': exec_res
         })
 
 class TestExecFallback(unittest.TestCase):
     def test_successful_execution(self):
         """Test that exec_fallback is not called when execution succeeds"""
-        shared_storage = {}
+        shared = {}
         node = FallbackNode(should_fail=False)
-        result = node.run(shared_storage)
+        result = node.run(shared)
         
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['attempts'], 1)
-        self.assertEqual(shared_storage['results'][0]['result'], "success")
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['attempts'], 1)
+        self.assertEqual(shared['results'][0]['result'], "success")
 
     def test_fallback_after_failure(self):
         """Test that exec_fallback is called after all retries are exhausted"""
-        shared_storage = {}
+        shared = {}
         node = FallbackNode(should_fail=True, max_retries=2)
-        result = node.run(shared_storage)
+        result = node.run(shared)
         
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['attempts'], 2)
-        self.assertEqual(shared_storage['results'][0]['result'], "fallback")
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['attempts'], 2)
+        self.assertEqual(shared['results'][0]['result'], "fallback")
 
     def test_fallback_in_flow(self):
         """Test that fallback works within a Flow"""
         class ResultNode(Node):
-            def prep(self, shared_storage):
-                return shared_storage.get('results', [])
+            def prep(self, shared):
+                return shared.get('results', [])
                 
-            def exec(self, prep_result):
-                return prep_result
+            def exec(self, prep_res):
+                return prep_res
                 
-            def post(self, shared_storage, prep_result, exec_result):
-                shared_storage['final_result'] = exec_result
+            def post(self, shared, prep_res, exec_res):
+                shared['final_result'] = exec_res
                 return None
         
-        shared_storage = {}
+        shared = {}
         fallback_node = FallbackNode(should_fail=True)
         result_node = ResultNode()
         fallback_node >> result_node
         
         flow = Flow(start=fallback_node)
-        flow.run(shared_storage)
+        flow.run(shared)
         
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['result'], "fallback")
-        self.assertEqual(shared_storage['final_result'], [{'attempts': 1, 'result': 'fallback'}] )
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['result'], "fallback")
+        self.assertEqual(shared['final_result'], [{'attempts': 1, 'result': 'fallback'}] )
 
     def test_no_fallback_implementation(self):
         """Test that default fallback behavior raises the exception"""
         class NoFallbackNode(Node):
-            def prep(self, shared_storage):
-                if 'results' not in shared_storage:
-                    shared_storage['results'] = []
+            def prep(self, shared):
+                if 'results' not in shared:
+                    shared['results'] = []
                 return None
             
-            def exec(self, prep_result):
+            def exec(self, prep_res):
                 raise ValueError("Test error")
             
-            def post(self, shared_storage, prep_result, exec_result):
-                shared_storage['results'].append({'result': exec_result})
-                return exec_result
+            def post(self, shared, prep_res, exec_res):
+                shared['results'].append({'result': exec_res})
+                return exec_res
         
-        shared_storage = {}
+        shared = {}
         node = NoFallbackNode()
         with self.assertRaises(ValueError):
-            node.run(shared_storage)
+            node.run(shared)
 
     def test_retry_before_fallback(self):
         """Test that retries are attempted before calling fallback"""
-        shared_storage = {}
+        shared = {}
         node = FallbackNode(should_fail=True, max_retries=3)
-        node.run(shared_storage)
+        node.run(shared)
         
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['attempts'], 3)
-        self.assertEqual(shared_storage['results'][0]['result'], "fallback")
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['attempts'], 3)
+        self.assertEqual(shared['results'][0]['result'], "fallback")
 
 class TestAsyncExecFallback(unittest.TestCase):
     def setUp(self):
@@ -146,77 +147,77 @@ def tearDown(self):
     def test_async_successful_execution(self):
         """Test that async exec_fallback is not called when execution succeeds"""
         async def run_test():
-            shared_storage = {}
+            shared = {}
             node = AsyncFallbackNode(should_fail=False)
-            await node.run_async(shared_storage)
-            return shared_storage
+            await node.run_async(shared)
+            return shared
         
-        shared_storage = self.loop.run_until_complete(run_test())
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['attempts'], 1)
-        self.assertEqual(shared_storage['results'][0]['result'], "success")
+        shared = self.loop.run_until_complete(run_test())
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['attempts'], 1)
+        self.assertEqual(shared['results'][0]['result'], "success")
 
     def test_async_fallback_after_failure(self):
         """Test that async exec_fallback is called after all retries are exhausted"""
         async def run_test():
-            shared_storage = {}
+            shared = {}
             node = AsyncFallbackNode(should_fail=True, max_retries=2)
-            await node.run_async(shared_storage)
-            return shared_storage
+            await node.run_async(shared)
+            return shared
         
-        shared_storage = self.loop.run_until_complete(run_test())
+        shared = self.loop.run_until_complete(run_test())
         
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['attempts'], 2)
-        self.assertEqual(shared_storage['results'][0]['result'], "async_fallback")
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['attempts'], 2)
+        self.assertEqual(shared['results'][0]['result'], "async_fallback")
 
     def test_async_fallback_in_flow(self):
         """Test that async fallback works within an AsyncFlow"""
         class AsyncResultNode(AsyncNode):
-            async def prep_async(self, shared_storage):
-                return shared_storage['results'][-1]['result']  # Get last result
+            async def prep_async(self, shared):
+                return shared['results'][-1]['result']  # Get last result
                 
-            async def exec_async(self, prep_result):
-                return prep_result
+            async def exec_async(self, prep_res):
+                return prep_res
                 
-            async def post_async(self, shared_storage, prep_result, exec_result):
-                shared_storage['final_result'] = exec_result
+            async def post_async(self, shared, prep_res, exec_res):
+                shared['final_result'] = exec_res
                 return "done"
         
         async def run_test():
-            shared_storage = {}
+            shared = {}
             fallback_node = AsyncFallbackNode(should_fail=True)
             result_node = AsyncResultNode()
             fallback_node >> result_node
             
             flow = AsyncFlow(start=fallback_node)
-            await flow.run_async(shared_storage)
-            return shared_storage
+            await flow.run_async(shared)
+            return shared
         
-        shared_storage = self.loop.run_until_complete(run_test())
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['result'], "async_fallback")
-        self.assertEqual(shared_storage['final_result'], "async_fallback")
+        shared = self.loop.run_until_complete(run_test())
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['result'], "async_fallback")
+        self.assertEqual(shared['final_result'], "async_fallback")
 
     def test_async_no_fallback_implementation(self):
         """Test that default async fallback behavior raises the exception"""
         class NoFallbackAsyncNode(AsyncNode):
-            async def prep_async(self, shared_storage):
-                if 'results' not in shared_storage:
-                    shared_storage['results'] = []
+            async def prep_async(self, shared):
+                if 'results' not in shared:
+                    shared['results'] = []
                 return None
             
-            async def exec_async(self, prep_result):
+            async def exec_async(self, prep_res):
                 raise ValueError("Test async error")
             
-            async def post_async(self, shared_storage, prep_result, exec_result):
-                shared_storage['results'].append({'result': exec_result})
-                return exec_result
+            async def post_async(self, shared, prep_res, exec_res):
+                shared['results'].append({'result': exec_res})
+                return exec_res
         
         async def run_test():
-            shared_storage = {}
+            shared = {}
             node = NoFallbackAsyncNode()
-            await node.run_async(shared_storage)
+            await node.run_async(shared)
         
         with self.assertRaises(ValueError):
             self.loop.run_until_complete(run_test())
@@ -224,15 +225,15 @@ async def run_test():
     def test_async_retry_before_fallback(self):
         """Test that retries are attempted before calling async fallback"""
         async def run_test():
-            shared_storage = {}
+            shared = {}
             node = AsyncFallbackNode(should_fail=True, max_retries=3)
-            result = await node.run_async(shared_storage)
-            return result, shared_storage
+            result = await node.run_async(shared)
+            return result, shared
         
-        result, shared_storage = self.loop.run_until_complete(run_test())
-        self.assertEqual(len(shared_storage['results']), 1)
-        self.assertEqual(shared_storage['results'][0]['attempts'], 3)
-        self.assertEqual(shared_storage['results'][0]['result'], "async_fallback")
+        result, shared = self.loop.run_until_complete(run_test())
+        self.assertEqual(len(shared['results']), 1)
+        self.assertEqual(shared['results'][0]['attempts'], 3)
+        self.assertEqual(shared['results'][0]['result'], "async_fallback")
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/tests/test_flow_basic.py b/tests/test_flow_basic.py
index e7b6dddc..6d93c9f9 100644
--- a/tests/test_flow_basic.py
+++ b/tests/test_flow_basic.py
@@ -1,11 +1,11 @@
 # tests/test_flow_basic.py
-import unittest
 import sys
-from pathlib import Path
+import unittest
 import warnings
+from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import Node, Flow
+from pocketflow import Flow, Node
 
 # --- Node Definitions ---
 # Nodes intended for default transitions (>>) should NOT return a specific
@@ -16,33 +16,33 @@ class NumberNode(Node):
     def __init__(self, number):
         super().__init__()
         self.number = number
-    def prep(self, shared_storage):
-        shared_storage['current'] = self.number
+    def prep(self, shared):
+        shared['current'] = self.number
     # post implicitly returns None - used for default transition
 
 class AddNode(Node):
     def __init__(self, number):
         super().__init__()
         self.number = number
-    def prep(self, shared_storage):
-        shared_storage['current'] += self.number
+    def prep(self, shared):
+        shared['current'] += self.number
     # post implicitly returns None - used for default transition
 
 class MultiplyNode(Node):
     def __init__(self, number):
         super().__init__()
         self.number = number
-    def prep(self, shared_storage):
-        shared_storage['current'] *= self.number
+    def prep(self, shared):
+        shared['current'] *= self.number
     # post implicitly returns None - used for default transition
 
 class CheckPositiveNode(Node):
    # This node IS designed for conditional branching
-   def prep(self, shared_storage):
+   def prep(self, shared):
        pass
-   def post(self, shared_storage, prep_result, proc_result):
+   def post(self, shared, prep_res, exec_res):
         # MUST return the specific action string for branching
-        if shared_storage['current'] >= 0:
+        if shared['current'] >= 0:
             return 'positive'
         else:
             return 'negative'
@@ -56,7 +56,7 @@ class EndSignalNode(Node):
     def __init__(self, signal="finished"):
         super().__init__()
         self.signal = signal
-    def post(self, shared_storage, prep_result, exec_result):
+    def post(self, shared, prep_res, exec_res):
         return self.signal # Return a specific signal
 
 # --- Test Class ---
@@ -64,30 +64,30 @@ class TestFlowBasic(unittest.TestCase):
 
     def test_start_method_initialization(self):
         """Test initializing flow with start() after creation."""
-        shared_storage = {}
+        shared = {}
         n1 = NumberNode(5)
         pipeline = Flow()
         pipeline.start(n1)
-        last_action = pipeline.run(shared_storage)
-        self.assertEqual(shared_storage['current'], 5)
+        last_action = pipeline.run(shared=shared)
+        self.assertEqual(shared['current'], 5)
         # NumberNode.post returns None (default)
         self.assertIsNone(last_action)
 
     def test_start_method_chaining(self):
         """Test fluent chaining using start().next()..."""
-        shared_storage = {}
+        shared = {}
         pipeline = Flow()
         # Chain: NumberNode -> AddNode -> MultiplyNode
         # All use default transitions (post returns None)
         pipeline.start(NumberNode(5)).next(AddNode(3)).next(MultiplyNode(2))
-        last_action = pipeline.run(shared_storage)
-        self.assertEqual(shared_storage['current'], 16)
+        last_action = pipeline.run(shared=shared)
+        self.assertEqual(shared['current'], 16)
         # Last node (MultiplyNode) post returns None
         self.assertIsNone(last_action)
 
     def test_sequence_with_rshift(self):
         """Test a simple linear pipeline using >>"""
-        shared_storage = {}
+        shared = {}
         n1 = NumberNode(5)
         n2 = AddNode(3)
         n3 = MultiplyNode(2)
@@ -96,14 +96,14 @@ def test_sequence_with_rshift(self):
         # All default transitions (post returns None)
         pipeline.start(n1) >> n2 >> n3
 
-        last_action = pipeline.run(shared_storage)
-        self.assertEqual(shared_storage['current'], 16)
+        last_action = pipeline.run(shared)
+        self.assertEqual(shared['current'], 16)
         # Last node (n3: MultiplyNode) post returns None
         self.assertIsNone(last_action)
 
     def test_branching_positive(self):
         """Test positive branch: CheckPositiveNode returns 'positive'"""
-        shared_storage = {}
+        shared = {}
         start_node = NumberNode(5)    # post -> None
         check_node = CheckPositiveNode() # post -> 'positive' or 'negative'
         add_if_positive = AddNode(10) # post -> None
@@ -116,14 +116,14 @@ def test_branching_positive(self):
         check_node - "negative" >> add_if_negative
 
         # Execution: start_node -> check_node -> add_if_positive
-        last_action = pipeline.run(shared_storage)
-        self.assertEqual(shared_storage['current'], 15) # 5 + 10
+        last_action = pipeline.run(shared)
+        self.assertEqual(shared['current'], 15) # 5 + 10
         # Last node executed was add_if_positive, its post returns None
         self.assertIsNone(last_action)
 
     def test_branching_negative(self):
         """Test negative branch: CheckPositiveNode returns 'negative'"""
-        shared_storage = {}
+        shared = {}
         start_node = NumberNode(-5)   # post -> None
         check_node = CheckPositiveNode() # post -> 'positive' or 'negative'
         add_if_positive = AddNode(10) # post -> None (won't run)
@@ -135,14 +135,14 @@ def test_branching_negative(self):
         check_node - "negative" >> add_if_negative
 
         # Execution: start_node -> check_node -> add_if_negative
-        last_action = pipeline.run(shared_storage)
-        self.assertEqual(shared_storage['current'], -25) # -5 + -20
+        last_action = pipeline.run(shared)
+        self.assertEqual(shared['current'], -25) # -5 + -20
         # Last node executed was add_if_negative, its post returns None
         self.assertIsNone(last_action)
 
     def test_cycle_until_negative_ends_with_signal(self):
         """Test cycle, ending on a node that returns a signal"""
-        shared_storage = {}
+        shared = {}
         n1 = NumberNode(10)           # post -> None
         check = CheckPositiveNode()   # post -> 'positive' or 'negative'
         subtract3 = AddNode(-3)       # post -> None
@@ -157,17 +157,17 @@ def test_cycle_until_negative_ends_with_signal(self):
         subtract3 >> check
 
         # Execution: n1->check->sub3->check->sub3->check->sub3->check->sub3->check->end_node
-        last_action = pipeline.run(shared_storage)
-        self.assertEqual(shared_storage['current'], -2) # 10 -> 7 -> 4 -> 1 -> -2
+        last_action = pipeline.run(shared=shared)
+        self.assertEqual(shared['current'], -2) # 10 -> 7 -> 4 -> 1 -> -2
         # Last node executed was end_node, its post returns "cycle_done"
         self.assertEqual(last_action, "cycle_done")
 
     def test_flow_ends_warning_default_missing(self):
         """Test warning when default transition is needed but not found"""
-        shared_storage = {}
+        shared = {}
         # Node that returns a specific action from post
         class ActionNode(Node):
-            def post(self, *args): return "specific_action"
+            def post(self, shared, prep_res, exec_res): return "specific_action"
         start_node = ActionNode()
         next_node = NoOpNode()
 
@@ -177,13 +177,13 @@ def post(self, *args): return "specific_action"
         start_node - "specific_action" >> next_node
 
         # Make start_node return None instead, triggering default search
-        start_node.post = lambda *args: None
+        start_node.post = lambda shared, prep_res, exec_res: None
 
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter("always")
             # Run flow. start_node runs, post returns None.
             # Flow looks for "default", but only "specific_action" exists.
-            last_action = pipeline.run(shared_storage)
+            last_action = pipeline.run(shared=shared)
 
             self.assertEqual(len(w), 1)
             self.assertTrue(issubclass(w[-1].category, UserWarning))
@@ -194,10 +194,10 @@ def post(self, *args): return "specific_action"
 
     def test_flow_ends_warning_specific_missing(self):
         """Test warning when specific action is returned but not found"""
-        shared_storage = {}
+        shared = {}
         # Node that returns a specific action from post
         class ActionNode(Node):
-            def post(self, *args): return "specific_action"
+            def post(self, *args, **kwargs): return "specific_action"
         start_node = ActionNode()
         next_node = NoOpNode()
 
@@ -210,7 +210,7 @@ def post(self, *args): return "specific_action"
             warnings.simplefilter("always")
             # Run flow. start_node runs, post returns "specific_action".
             # Flow looks for "specific_action", but only "default" exists.
-            last_action = pipeline.run(shared_storage)
+            last_action = pipeline.run(shared)
 
             self.assertEqual(len(w), 1)
             self.assertTrue(issubclass(w[-1].category, UserWarning))
diff --git a/tests/test_flow_composition.py b/tests/test_flow_composition.py
index 544bd08e..bcc6324e 100644
--- a/tests/test_flow_composition.py
+++ b/tests/test_flow_composition.py
@@ -1,35 +1,35 @@
 # tests/test_flow_composition.py
-import unittest
-import asyncio # Keep import, might be needed if other tests use it indirectly
 import sys
+import unittest
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
-from pocketflow import Node, Flow
+from pocketflow import Flow, Node
+
 
 # --- Existing Nodes ---
 class NumberNode(Node):
     def __init__(self, number):
         super().__init__()
         self.number = number
-    def prep(self, shared_storage):
-        shared_storage['current'] = self.number
+    def prep(self, shared):
+        shared['current'] = self.number
     # post implicitly returns None
 
 class AddNode(Node):
     def __init__(self, number):
         super().__init__()
         self.number = number
-    def prep(self, shared_storage):
-        shared_storage['current'] += self.number
+    def prep(self, shared):
+        shared['current'] += self.number
     # post implicitly returns None
 
 class MultiplyNode(Node):
     def __init__(self, number):
         super().__init__()
         self.number = number
-    def prep(self, shared_storage):
-        shared_storage['current'] *= self.number
+    def prep(self, shared):
+        shared['current'] *= self.number
     # post implicitly returns None
 
 # --- New Nodes for Action Propagation Test ---
@@ -39,9 +39,9 @@ def __init__(self, signal="default_signal"):
         super().__init__()
         self.signal = signal
     # No prep needed usually if just signaling
-    def post(self, shared_storage, prep_result, exec_result):
+    def post(self, shared, prep_res, exec_res):
         # Store the signal in shared storage for verification
-        shared_storage['last_signal_emitted'] = self.signal
+        shared['last_signal_emitted'] = self.signal
         return self.signal # Return the specific action string
 
 class PathNode(Node):
@@ -49,8 +49,8 @@ class PathNode(Node):
     def __init__(self, path_id):
         super().__init__()
         self.path_id = path_id
-    def prep(self, shared_storage):
-        shared_storage['path_taken'] = self.path_id
+    def prep(self, shared):
+        shared['path_taken'] = self.path_id
     # post implicitly returns None
 
 # --- Test Class ---
@@ -62,15 +62,15 @@ def test_flow_as_node(self):
         1) Create a Flow (f1) starting with NumberNode(5), then AddNode(10), then MultiplyNode(2).
         2) Create a second Flow (f2) whose start is f1.
         3) Create a wrapper Flow (f3) that contains f2 to ensure proper execution.
-        Expected final result in shared_storage['current']: (5 + 10) * 2 = 30.
+        Expected final result in shared['current']: (5 + 10) * 2 = 30.
         """
-        shared_storage = {}
+        shared = {}
         f1 = Flow(start=NumberNode(5))
         f1 >> AddNode(10) >> MultiplyNode(2)
         f2 = Flow(start=f1)
         f3 = Flow(start=f2)
-        f3.run(shared_storage)
-        self.assertEqual(shared_storage['current'], 30)
+        f3.run(shared)
+        self.assertEqual(shared['current'], 30)
 
     def test_nested_flow(self):
         """
@@ -80,14 +80,14 @@ def test_nested_flow(self):
         wrapper_flow: contains middle_flow to ensure proper execution
         Expected final result: (5 + 3) * 4 = 32.
         """
-        shared_storage = {}
+        shared = {}
         inner_flow = Flow(start=NumberNode(5))
         inner_flow >> AddNode(3)
         middle_flow = Flow(start=inner_flow)
         middle_flow >> MultiplyNode(4)
         wrapper_flow = Flow(start=middle_flow)
-        wrapper_flow.run(shared_storage)
-        self.assertEqual(shared_storage['current'], 32)
+        wrapper_flow.run(shared)
+        self.assertEqual(shared['current'], 32)
 
     def test_flow_chaining_flows(self):
         """
@@ -97,22 +97,22 @@ def test_flow_chaining_flows(self):
         wrapper_flow: contains both flow1 and flow2 to ensure proper execution
         Expected final result: (10 + 10) * 2 = 40.
         """
-        shared_storage = {}
+        shared = {}
         numbernode = NumberNode(10)
         numbernode >> AddNode(10)
         flow1 = Flow(start=numbernode)
         flow2 = Flow(start=MultiplyNode(2))
         flow1 >> flow2 # Default transition based on flow1 returning None
         wrapper_flow = Flow(start=flow1)
-        wrapper_flow.run(shared_storage)
-        self.assertEqual(shared_storage['current'], 40)
+        wrapper_flow.run(shared)
+        self.assertEqual(shared['current'], 40)
 
     def test_composition_with_action_propagation(self):
         """
         Test that an outer flow can branch based on the action returned
         by the last node's post() within an inner flow.
         """
-        shared_storage = {}
+        shared = {}
 
         # 1. Define an inner flow that ends with a node returning a specific action
         inner_start_node = NumberNode(100)       # current = 100, post -> None
@@ -135,14 +135,14 @@ def test_composition_with_action_propagation(self):
 
         # 5. Run the outer flow and capture the last action
         # Execution: inner_start -> inner_end -> path_b
-        last_action_outer = outer_flow.run(shared_storage)
+        last_action_outer = outer_flow.run(shared)
 
         # 6. Assert the results
         # Check state after inner flow execution
-        self.assertEqual(shared_storage.get('current'), 100)
-        self.assertEqual(shared_storage.get('last_signal_emitted'), "inner_done")
+        self.assertEqual(shared.get('current'), 100)
+        self.assertEqual(shared.get('last_signal_emitted'), "inner_done")
         # Check that the correct outer path was taken
-        self.assertEqual(shared_storage.get('path_taken'), "B")
+        self.assertEqual(shared.get('path_taken'), "B")
         # Check the action returned by the outer flow. The last node executed was
         # path_b_node, which returns None from its post method.
         self.assertIsNone(last_action_outer)