From 442fd269fa6a826cf584dc17a96488b03ddafce1 Mon Sep 17 00:00:00 2001
From: zachary62 <zhuang333@wisc.edu>
Date: Sun, 25 May 2025 19:39:17 -0400
Subject: [PATCH] update streamlit nodes/flow

---
 cookbook/pocketflow-streamlit-hitl/flow.py  | 25 ++-------
 cookbook/pocketflow-streamlit-hitl/nodes.py | 62 ++++-----------------
 2 files changed, 16 insertions(+), 71 deletions(-)

diff --git a/cookbook/pocketflow-streamlit-hitl/flow.py b/cookbook/pocketflow-streamlit-hitl/flow.py
index d204c28..ada5a66 100644
--- a/cookbook/pocketflow-streamlit-hitl/flow.py
+++ b/cookbook/pocketflow-streamlit-hitl/flow.py
@@ -1,24 +1,9 @@
 from pocketflow import Flow
-from nodes import InitialInputNode, ProcessDataNode, PrepareFinalResultNode
+from nodes import GenerateImageNode
 
-def create_initial_processing_flow():
-    """Creates a flow for the initial data processing stage."""
-    initial_input_node = InitialInputNode()
-    process_data_node = ProcessDataNode()
+def create_generation_flow():
+    """Creates a flow for image generation (initial or regeneration)."""
+    generate_image_node = GenerateImageNode()
+    return Flow(start=generate_image_node)
 
-    # Define transitions: Input -> Process
-    initial_input_node >> process_data_node
 
-    # Create the Flow, starting with the input node
-    flow = Flow(start=initial_input_node)
-    print("Initial processing flow created.")
-    return flow
-
-def create_finalization_flow():
-    """Creates a flow to finalize the result after approval."""
-    prepare_final_result_node = PrepareFinalResultNode()
-
-    # This flow only has one node
-    flow = Flow(start=prepare_final_result_node)
-    print("Finalization flow created.")
-    return flow
diff --git a/cookbook/pocketflow-streamlit-hitl/nodes.py b/cookbook/pocketflow-streamlit-hitl/nodes.py
index 5a265be..e530696 100644
--- a/cookbook/pocketflow-streamlit-hitl/nodes.py
+++ b/cookbook/pocketflow-streamlit-hitl/nodes.py
@@ -1,57 +1,17 @@
 from pocketflow import Node
-from utils.process_task import process_task
+from utils.generate_image import generate_image
 
-class InitialInputNode(Node):
-    """Reads the initial task input from shared_data."""
+class GenerateImageNode(Node):
+    """Generates image from text prompt using OpenAI API."""
+    
     def prep(self, shared):
-        print("InitialInputNode: Prep")
-        return shared.get("task_input", "Default Task Input")
+        return shared.get("task_input", "")
 
-    def exec(self, prep_res):
-        print(f"InitialInputNode: Executing with input: '{prep_res[:50]}...'")
-        # No real computation needed here, just passing the input along
-        return prep_res
+    def exec(self, prompt):
+        return generate_image(prompt)
 
     def post(self, shared, prep_res, exec_res):
-        # Ensure the input used is stored, although it might already be there
-        shared["input_used_by_process"] = exec_res
-        print(f"InitialInputNode: Post - Stored input '{exec_res[:50]}...' in shared_data.")
-        return "default" # Proceed to next node in the flow
-
-class ProcessDataNode(Node):
-    """Processes the data using the utility function."""
-    def prep(self, shared):
-        task_input = shared.get("input_used_by_process", "No input found")
-        print(f"ProcessDataNode: Prep - Input: '{task_input[:50]}...'")
-        return task_input
-
-    def exec(self, prep_res):
-        print("ProcessDataNode: Exec - Calling process_task utility")
-        # Call the actual processing logic
-        processed_output = process_task(prep_res)
-        return processed_output
-
-    def post(self, shared, prep_res, exec_res):
-        # Store the result for review
-        shared["processed_output"] = exec_res
-        print(f"ProcessDataNode: Post - Stored processed output: '{str(exec_res)[:50]}...'")
-        # This node ends the initial processing subflow
-        return None
-
-class PrepareFinalResultNode(Node):
-    """Takes the approved processed output and sets it as the final result."""
-    def prep(self, shared):
-        approved_output = shared.get("processed_output", "No processed output found")
-        print(f"PrepareFinalResultNode: Prep - Approved output: '{str(approved_output)[:50]}...'")
-        return approved_output
-
-    def exec(self, prep_res):
-        print("PrepareFinalResultNode: Exec - Finalizing result.")
-        # Could potentially do final formatting here if needed
-        return prep_res
-
-    def post(self, shared, prep_res, exec_res):
-        shared["final_result"] = exec_res
-        print(f"PrepareFinalResultNode: Post - Stored final result: '{str(exec_res)[:50]}...'")
-        # This node ends the finalization subflow
-        return None
\ No newline at end of file
+        shared["input_used_by_process"] = prep_res
+        shared["generated_image"] = exec_res
+        shared["stage"] = "user_feedback"
+        return "default"
\ No newline at end of file