125 lines
2.8 KiB
Python
125 lines
2.8 KiB
Python
"""
|
|
PocketFlow Node Template
|
|
|
|
Copy this template and customize for your needs
|
|
"""
|
|
|
|
from pocketflow import Node
|
|
# from utils.call_llm import call_llm # Uncomment if using LLM
|
|
|
|
|
|
class TemplateNode(Node):
|
|
"""
|
|
Brief description of what this node does
|
|
|
|
Shared Store Schema:
|
|
Input:
|
|
- key1 (type): description
|
|
- key2 (type): description
|
|
|
|
Output:
|
|
- result_key (type): description
|
|
|
|
Actions:
|
|
- "default": Normal flow
|
|
- "error": If something goes wrong
|
|
- "retry": If needs retry
|
|
"""
|
|
|
|
def prep(self, shared):
|
|
"""
|
|
Prepare data from shared store
|
|
|
|
Args:
|
|
shared (dict): Shared data store
|
|
|
|
Returns:
|
|
Any: Data to pass to exec()
|
|
"""
|
|
# TODO: Get data from shared store
|
|
input_data = shared.get("input_key")
|
|
|
|
# Optional: Add validation
|
|
if not input_data:
|
|
raise ValueError("Missing required input")
|
|
|
|
return input_data
|
|
|
|
def exec(self, prep_res):
|
|
"""
|
|
Execute the main logic (can fail and retry)
|
|
|
|
Args:
|
|
prep_res: Data from prep()
|
|
|
|
Returns:
|
|
Any: Result to pass to post()
|
|
"""
|
|
# TODO: Implement your logic here
|
|
|
|
# Example: Call LLM
|
|
# result = call_llm(f"Process: {prep_res}")
|
|
|
|
# Example: Process data
|
|
result = f"Processed: {prep_res}"
|
|
|
|
return result
|
|
|
|
def post(self, shared, prep_res, exec_res):
|
|
"""
|
|
Save results and return action
|
|
|
|
Args:
|
|
shared (dict): Shared data store
|
|
prep_res: Original data from prep()
|
|
exec_res: Result from exec()
|
|
|
|
Returns:
|
|
str: Action name for flow control
|
|
"""
|
|
# TODO: Save results to shared store
|
|
shared["result_key"] = exec_res
|
|
|
|
# Optional: Conditional actions
|
|
# if some_condition:
|
|
# return "special_action"
|
|
|
|
return "default"
|
|
|
|
def exec_fallback(self, prep_res, exc):
|
|
"""
|
|
Optional: Handle errors gracefully
|
|
|
|
Args:
|
|
prep_res: Data from prep()
|
|
exc: The exception that occurred
|
|
|
|
Returns:
|
|
Any: Fallback result (passed to post as exec_res)
|
|
"""
|
|
# TODO: Implement fallback logic
|
|
print(f"Error occurred: {exc}")
|
|
|
|
# Option 1: Re-raise the exception
|
|
# raise exc
|
|
|
|
# Option 2: Return fallback value
|
|
return "Fallback result"
|
|
|
|
|
|
# Example usage
|
|
if __name__ == "__main__":
|
|
# Create node with retry settings
|
|
node = TemplateNode(max_retries=3, wait=5)
|
|
|
|
# Create shared store
|
|
shared = {
|
|
"input_key": "test input"
|
|
}
|
|
|
|
# Run node
|
|
action = node.run(shared)
|
|
|
|
print(f"Action: {action}")
|
|
print(f"Result: {shared.get('result_key')}")
|