136 lines
3.7 KiB
Python
136 lines
3.7 KiB
Python
import unittest
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.append(str(Path(__file__).parent.parent))
|
|
from minillmflow import Node, Flow
|
|
|
|
class NumberNode(Node):
|
|
def __init__(self, number):
|
|
super().__init__()
|
|
self.number = number
|
|
|
|
def process(self, shared_storage, prep_result):
|
|
shared_storage['current'] = self.number
|
|
|
|
class AddNode(Node):
|
|
def __init__(self, number):
|
|
super().__init__()
|
|
self.number = number
|
|
|
|
def process(self, shared_storage, prep_result):
|
|
shared_storage['current'] += self.number
|
|
|
|
class MultiplyNode(Node):
|
|
def __init__(self, number):
|
|
super().__init__()
|
|
self.number = number
|
|
|
|
def process(self, shared_storage, prep_result):
|
|
shared_storage['current'] *= self.number
|
|
|
|
|
|
class TestFlowComposition(unittest.TestCase):
|
|
|
|
def test_flow_as_node(self):
|
|
"""
|
|
Demonstrates that a Flow can itself be chained like a Node.
|
|
We create a flow (f1) that starts with NumberNode(5) -> AddNode(10).
|
|
Then we chain f1 >> MultiplyNode(2).
|
|
|
|
Expected result after running from f1:
|
|
start = 5
|
|
5 + 10 = 15
|
|
15 * 2 = 30
|
|
"""
|
|
shared_storage = {}
|
|
|
|
# Inner flow f1
|
|
f1 = Flow(start_node=NumberNode(5))
|
|
f1 >> AddNode(10)
|
|
|
|
# Then chain a node after the flow
|
|
f1 >> MultiplyNode(2)
|
|
|
|
# Run from f1
|
|
f1.run(shared_storage)
|
|
|
|
self.assertEqual(shared_storage['current'], 30)
|
|
|
|
def test_nested_flow(self):
|
|
"""
|
|
Demonstrates embedding one Flow inside another Flow.
|
|
inner_flow: NumberNode(5) -> AddNode(3)
|
|
outer_flow: starts with inner_flow -> MultiplyNode(4)
|
|
|
|
Expected result:
|
|
(5 + 3) * 4 = 32
|
|
"""
|
|
shared_storage = {}
|
|
|
|
# Define an inner flow
|
|
inner_flow = Flow(start_node=NumberNode(5))
|
|
inner_flow >> AddNode(3)
|
|
|
|
# Define an outer flow, whose start node is inner_flow
|
|
outer_flow = Flow(start_node=inner_flow)
|
|
outer_flow >> MultiplyNode(4)
|
|
|
|
# Run outer_flow
|
|
outer_flow.run(shared_storage)
|
|
|
|
self.assertEqual(shared_storage['current'], 32) # (5+3)*4=32
|
|
|
|
def test_flow_chaining_flows(self):
|
|
"""
|
|
Demonstrates chaining one flow to another flow.
|
|
flow1: NumberNode(10) -> AddNode(10) # final shared_storage['current'] = 20
|
|
flow2: MultiplyNode(2) # final shared_storage['current'] = 40
|
|
|
|
flow1 >> flow2 means once flow1 finishes, flow2 starts.
|
|
|
|
Expected result: (10 + 10) * 2 = 40
|
|
"""
|
|
shared_storage = {}
|
|
|
|
# flow1
|
|
flow1 = Flow(start_node=NumberNode(10))
|
|
flow1 >> AddNode(10)
|
|
|
|
# flow2
|
|
flow2 = Flow(start_node=MultiplyNode(2))
|
|
|
|
# Chain them: flow1 >> flow2
|
|
flow1 >> flow2
|
|
|
|
# Start running from flow1
|
|
flow1.run(shared_storage)
|
|
|
|
self.assertEqual(shared_storage['current'], 40)
|
|
|
|
def test_flow_with_parameters(self):
|
|
"""
|
|
Demonstrates passing parameters into a Flow (and retrieved by a Node).
|
|
"""
|
|
|
|
class ParamNode(Node):
|
|
def process(self, shared_storage, prep_result):
|
|
# Reads 'level' from the node's (or flow's) parameters
|
|
shared_storage['param'] = self.parameters.get('level', 'no param')
|
|
|
|
shared_storage = {}
|
|
|
|
# Create a flow with a ParamNode
|
|
f = Flow(start_node=ParamNode())
|
|
# Set parameters on the flow
|
|
f.parameters = {'level': 'Level 1'}
|
|
|
|
f.run(shared_storage)
|
|
|
|
self.assertEqual(shared_storage['param'], 'Level 1')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|