Update readme

README.md

@@ -7,36 +7,37 @@
 ## Usage
 
 ```python
-from permchain import InMemoryPubSubConnection, PubSub, Topic
+from permchain import Pregel, channels
 
-topic_one = Topic("one")
-chain_one = Topic.IN.subscribe() | (lambda x: x + 'b') | topic_one.publish()
-chain_two = topic_one.subscribe() | (lambda x: x + 'c') | Topic.OUT.publish()
+value = channels.LastValue[str]("value")
 
-conn = InMemoryPubSubConnection()
-pubsub = PubSub(chain_one, chain_two, connection=conn)
+grow_value = (
+    Pregel.subscribe_to(value)
+    | (lambda x: x + x)
+    | Pregel.send_to({value: lambda x: x if len(x) < 10 else None})
+)
 
-assert pubsub.invoke('a') == 'abc'
+pubsub = Pregel(grow_value, input=value, output=value)
+
+assert pubsub.invoke("a") == "aaaaaaaa"
 ```
 
-Check `tests` and `examples` for more examples.
+Check `examples` for more examples.
 
 ## Near-term Roadmap
 
-- [x] Add initial retry support (pending changes in `langchain`)
-- [x] Implement OUT as regular topic
-- [x] Implement IN as regular topic
-- [x] Add Connection.peek() to monitor past messages from all topics
-- [x] Enable resuming PubSub from the "middle" of the computation
-- [x] Add test for .peek()
-- [x] Add "wait until topic X is done" pattern, aka. `Topic.join()`
-- [ ] Move tracking of inflight processes/messages to Connection
-  - [ ] Use this to build retry mechanism, where any inflight messages are moved back to the respective topics when restarting
-  - [ ] But this would require being able to replay a message for a single listener only, which maybe requires a larger redesign of PubSub<>Connection contract than what I wanted to do here
-- [ ] Detect cycles (aka. infinite loops) and throw an error
-  - [ ] Allow user to catch that error (by subcribing to an error topic?)
-- [ ] Add example for "human in the loop" pattern, one of the two below
-  - [ ] Example with one permchain, which runs until it produces either 1. request for input or 2. output. The consumer code then gets the needed info, and restarts the permchain with answer, and same state id
-  - [ ] Allow interrupting execution by breaking out of the iterator returned by .stream()
-    - [ ] Build example showing a simple "human in the loop" pattern using this, ie. if a certain message asking for input is published the consumer of the iterator breaks out, does something and then restarts it
-- [ ] Add Redis-backed Connection implementation
+- [ ] Iterate on API
+  - [ ] do we want api to receive output from multiple channels in invoke()
+  - [ ] do we want api to send input to multiple channels in invoke()
+- [ ] Implement checkpointing
+  - [ ] Save checkpoints at end of each step
+  - [ ] Load checkpoint at start of invocation
+  - [ ] API to specify storage backend and save key
+- [ ] Add more examples
+  - [ ] human in the loop
+  - [ ] combine documents
+  - [ ] agent executor
+  - [ ] run over dataset
+- [ ] Fault tolerance
+  - [ ] Retry individual processes in a step
+  - [ ] Retry entire step?

examples/readme.py

@@ -0,0 +1,13 @@
+from permchain import Pregel, channels
+
+value = channels.LastValue[str]("value")
+
+grow_value = (
+    Pregel.subscribe_to(value)
+    | (lambda x: x + x)
+    | Pregel.send_to({value: lambda x: x if len(x) < 10 else None})
+)
+
+pubsub = Pregel(grow_value, input=value, output=value)
+
+assert pubsub.invoke("a") == "aaaaaaaa"

permchain/pregel.py

@@ -588,7 +588,3 @@ def _apply_writes_and_prepare_next_tasks(
                 tasks.append((proc, val))
 
     return tasks
-
-
-# TODO do we want api to subscribe to all channels?
-# Do we want api to send input to multiple channels in invoke()

tests/test_pregel.py

@@ -1,5 +1,4 @@
 import time
-from uuid import uuid4
 
 import pytest
 from pytest_mock import MockerFixture
@@ -10,28 +9,28 @@ from permchain.pregel import PregelInvoke
 
 
 def test_invoke_single_process_in_out(mocker: MockerFixture):
-    input = channels.LastValue[int]("input")
-    output = channels.LastValue[int]("output")
+    input_chan = channels.LastValue[int]("input")
+    output_chan = channels.LastValue[int]("output")
 
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output)
+    chain = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output_chan)
 
-    pubsub = Pregel(chain, input=input, output=output)
+    pubsub = Pregel(chain, input=input_chan, output=output_chan)
 
     # Then invoke pubsub
     assert pubsub.invoke(2) == 3
 
 
 def test_invoke_two_processes_in_out(mocker: MockerFixture):
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     inbox = channels.Inbox[int]("inbox")
 
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox)
+    chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox)
     chain_two = Pregel.subscribe_to_each(inbox) | add_one | Pregel.send_to(output)
 
-    pubsub = Pregel(chain_one, chain_two, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
 
     # Then invoke pubsub
     assert pubsub.invoke(2) == 4
@@ -42,13 +41,16 @@ def test_batch_two_processes_in_out(mocker: MockerFixture):
         time.sleep(inp / 10)
         return inp + 1
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     one = channels.LastValue[int]("one")
-    chain_one = Pregel.subscribe_to(input) | add_one_with_delay | Pregel.send_to(one)
+
+    chain_one = (
+        Pregel.subscribe_to(input_chan) | add_one_with_delay | Pregel.send_to(one)
+    )
     chain_two = Pregel.subscribe_to(one) | add_one_with_delay | Pregel.send_to(output)
 
-    pubsub = Pregel(chain_one, chain_two, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
 
     # Then invoke pubsub
     assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
@@ -57,13 +59,13 @@ def test_batch_two_processes_in_out(mocker: MockerFixture):
 def test_invoke_many_processes_in_out(mocker: MockerFixture):
     test_size = 100
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     topics: list[channels.Channel] = [channels.LastValue[int]("zero")]
 
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     chains: list[PregelInvoke] = [
-        Pregel.subscribe_to(input) | add_one | Pregel.send_to(topics[0])
+        Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0])
     ]
     for i in range(test_size - 2):
         topics.append(channels.LastValue[int](str(i)))
@@ -72,7 +74,7 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture):
         )
     chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output))
 
-    pubsub = Pregel(*chains, input=input, output=output)
+    pubsub = Pregel(*chains, input=input_chan, output=output)
 
     for _ in range(10):
         # Then invoke pubsub
@@ -82,13 +84,13 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture):
 def test_batch_many_processes_in_out(mocker: MockerFixture):
     test_size = 100
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     topics: list[channels.Channel] = [channels.LastValue[int]("zero")]
 
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     chains: list[PregelInvoke] = [
-        Pregel.subscribe_to(input) | add_one | Pregel.send_to(topics[0])
+        Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0])
     ]
     for i in range(test_size - 2):
         topics.append(channels.LastValue[int](str(i)))
@@ -97,7 +99,7 @@ def test_batch_many_processes_in_out(mocker: MockerFixture):
         )
     chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output))
 
-    pubsub = Pregel(*chains, input=input, output=output)
+    pubsub = Pregel(*chains, input=input_chan, output=output)
 
     for _ in range(10):
         # Then invoke pubsub
@@ -113,13 +115,13 @@ def test_batch_many_processes_in_out(mocker: MockerFixture):
 def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture):
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
 
-    chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output)
-    chain_two = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output)
+    chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
+    chain_two = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
 
-    pubsub = Pregel(chain_one, chain_two, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
 
     with pytest.raises(channels.InvalidUpdateError):
         # LastValue channels can only be updated once per iteration
@@ -129,13 +131,13 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture):
 def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture):
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.Inbox[int]("output")
 
-    chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output)
-    chain_two = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output)
+    chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
+    chain_two = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
 
-    pubsub = Pregel(chain_one, chain_two, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
 
     # An Inbox channel accumulates updates into a sequence
     assert pubsub.invoke(2) == (3, 3)
@@ -145,15 +147,15 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture):
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     inbox = channels.Inbox[int]("inbox")
 
-    chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox)
-    chain_three = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox)
+    chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox)
+    chain_three = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox)
     chain_four = Pregel.subscribe_to(inbox) | add_10_each | Pregel.send_to(output)
 
-    pubsub = Pregel(chain_one, chain_three, chain_four, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_three, chain_four, input=input_chan, output=output)
 
     # Then invoke pubsub
     # We get a single array result as chain_four waits for all publishers to finish
@@ -166,12 +168,12 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
 
     inner_pubsub = Pregel(
-        Pregel.subscribe_to(input) | add_one | Pregel.send_to(output),
-        input=input,
+        Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output),
+        input=input_chan,
         output=output,
     )
 
@@ -179,7 +181,7 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
     outbox_one = channels.LastValue[int]("outbox_one")
 
     chain_one = (
-        Pregel.subscribe_to(input) | add_10_each | Pregel.send_to(inbox_one).map()
+        Pregel.subscribe_to(input_chan) | add_10_each | Pregel.send_to(inbox_one).map()
     )
     chain_two = (
         Pregel.subscribe_to(inbox_one)
@@ -189,7 +191,7 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
     )
     chain_three = Pregel.subscribe_to(outbox_one) | sum | Pregel.send_to(output)
 
-    pubsub = Pregel(chain_one, chain_two, chain_three, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, chain_three, input=input_chan, output=output)
 
     # Then invoke pubsub
     for _ in range(10):
@@ -199,12 +201,12 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
 def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture):
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     between = channels.LastValue[int]("between")
 
     chain_one = (
-        Pregel.subscribe_to(input)
+        Pregel.subscribe_to(input_chan)
         | add_one
         | Pregel.send_to(
             {output: RunnablePassthrough(), between: RunnablePassthrough()}
@@ -212,22 +214,22 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture):
     )
     chain_two = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output)
 
-    pubsub = Pregel(chain_one, chain_two, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
 
     # Then invoke pubsub
     assert [c for c in pubsub.stream(2)] == [3, 4]
 
 
 def test_invoke_two_processes_no_out(mocker: MockerFixture):
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     between = channels.LastValue[int]("between")
 
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(between)
+    chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(between)
     chain_two = Pregel.subscribe_to(between) | add_one
 
-    pubsub = Pregel(chain_one, chain_two, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
 
     # Then invoke pubsub
     # It finishes executing (once no more messages being published)
@@ -236,7 +238,7 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture):
 
 
 def test_invoke_two_processes_no_in(mocker: MockerFixture):
-    input = channels.LastValue[int]("input")
+    input_chan = channels.LastValue[int]("input")
     output = channels.LastValue[int]("output")
     between = channels.LastValue[int]("between")
 
@@ -244,7 +246,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture):
     chain_one = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output)
     chain_two = Pregel.subscribe_to(between) | add_one
 
-    pubsub = Pregel(chain_one, chain_two, input=input, output=output)
+    pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
 
     with pytest.raises(ValueError):
         assert pubsub.invoke(2) is None