[Brook] Added full examples for the streams API proposal.

tests/contrib/streams/examples/README.md

@@ -0,0 +1,15 @@
+# Examples for the Brook streams API proposal
+
+Each file is a complete program for one section of the proposal. `test_examples.py` runs each one end to end on the memory Store and on Redis.
+
+| Proposal section | File |
+|---|---|
+| A UI backend reads a Workflow's Stream | [`ui_backend_reads.py`](ui_backend_reads.py) |
+| A client resumes a scoped Stream from an offset | [`resume_from_offset.py`](resume_from_offset.py) |
+| A client creates a Stream and passes it to a Workflow | [`client_creates_stream.py`](client_creates_stream.py) |
+| An agent Workflow publishes its turns | [`agent_publishes_turns.py`](agent_publishes_turns.py) |
+| An Activity streams tokens to the Workflow's Stream | [`activity_streams_tokens.py`](activity_streams_tokens.py) |
+| A Workflow scopes a Stream to one turn and closes it | [`turn_scoped_stream.py`](turn_scoped_stream.py) |
+| Worker setup | [`worker_setup.py`](worker_setup.py) |
+
+The proposal's Nexus examples are a later stage, so they have no file here.

tests/contrib/streams/examples/activity_streams_tokens.py

@@ -0,0 +1,99 @@
+"""Brook API proposal, "An Activity streams tokens to the Workflow's Stream".
+
+An async Activity publishes each token with an awaited ``publish``, which returns the offset,
+and batches the rest with ``BufferedPublisher``. A sync Activity takes ``SyncActivityStream``,
+which works in a thread pool executor.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import uuid
+from dataclasses import dataclass
+from datetime import timedelta
+
+from temporalio import activity, workflow
+from temporalio.client import Client
+from temporalio.contrib import streams
+from temporalio.contrib.streams import workflow as stream_workflow
+from temporalio.worker import Worker
+
+
+@streams.defn(name="events")
+class Events:
+    tokens = streams.topic(str)
+
+
+@dataclass
+class CallModelInput:
+    prompt: str
+    stream: streams.ActivityStream[Events]
+
+
+@dataclass
+class SummarizeInput:
+    text: str
+    stream: streams.SyncActivityStream[Events]
+
+
+@activity.defn
+async def call_model(input: CallModelInput) -> int:
+    tokens = input.stream.topic(Events.tokens)
+    words = input.prompt.split()
+    # Each awaited publish has written when it returns.
+    last = await tokens.publish(words[0])
+    # A fast model sends many small tokens, so the helper batches them.
+    async with streams.BufferedPublisher(
+        tokens, max_delay=timedelta(milliseconds=100)
+    ) as fast:
+        for word in words[1:]:
+            fast.publish(word)
+    return last
+
+
+@activity.defn
+def summarize(input: SummarizeInput) -> int:
+    return input.stream.topic(Events.tokens).publish(f"summary: {input.text}")
+
+
+@workflow.defn
+class Model:
+    @workflow.run
+    async def run(self, prompt: str) -> list[int]:
+        stream = await stream_workflow.create_stream(Events)
+        first = await workflow.execute_activity(
+            call_model,
+            CallModelInput(prompt, streams.ActivityStream.from_handle(stream)),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        summary = await workflow.execute_activity(
+            summarize,
+            SummarizeInput(prompt, streams.SyncActivityStream.from_handle(stream)),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        return [first, summary]
+
+
+async def main(
+    client: Client, task_queue: str
+) -> tuple[list[int], list[streams.Record[str]]]:
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        async with Worker(
+            client,
+            task_queue=task_queue,
+            workflows=[Model],
+            activities=[call_model, summarize],
+            activity_executor=executor,
+        ):
+            handle = await client.start_workflow(
+                Model.run,
+                "one two three",
+                id=f"model-{uuid.uuid4()}",
+                task_queue=task_queue,
+            )
+            offsets = await handle.result()
+    stream = streams.StreamClient(client).workflow_stream(handle, Events)
+    records = [
+        item async for item in stream.subscribe() if isinstance(item, streams.Record)
+    ]
+    return offsets, records

tests/contrib/streams/examples/agent_publishes_turns.py

@@ -0,0 +1,107 @@
+"""Brook API proposal, "An agent Workflow publishes its turns".
+
+The Workflow creates its Stream, publishes each turn with an awaited ``publish``, and passes the
+Stream to the model Activity in its one input object. ``close_with_run`` closes the Stream with
+the result of the run chain.
+"""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from datetime import timedelta
+
+from temporalio import activity, workflow
+from temporalio.client import Client
+from temporalio.contrib import streams
+from temporalio.contrib.streams import workflow as stream_workflow
+from temporalio.worker import Worker
+
+
+@dataclass
+class Turn:
+    kind: str
+    text: str = ""
+
+
+@streams.defn(name="events")
+class Events:
+    turns = streams.topic(Turn)
+    tokens = streams.topic(str)
+
+
+@dataclass
+class Reply:
+    text: str
+    tool_calls: list[str] = field(default_factory=list)
+
+
+@dataclass
+class CallModelInput:
+    messages: list[str]
+    stream: streams.ActivityStream[Events]
+
+
+@activity.defn
+async def call_model(input: CallModelInput) -> Reply:
+    await input.stream.topic(Events.tokens).publish(*input.messages[-1].split())
+    if len(input.messages) == 1:
+        return Reply("", tool_calls=["search"])
+    return Reply("the answer")
+
+
+@activity.defn
+async def run_tool(call: str) -> str:
+    return f"{call} result"
+
+
+@workflow.defn
+class AgentTurns:
+    @workflow.run
+    async def run(self, messages: list[str]) -> str:
+        stream = await stream_workflow.create_stream(
+            Events, retention=timedelta(days=1), auto_close=timedelta(hours=1)
+        )
+        turns = stream.topic(Events.turns)
+        while True:
+            # The publish goes in the run's buffer, so it and the model call start in one
+            # Workflow Task.
+            await turns.publish(Turn("model"))
+            reply = await workflow.execute_activity(
+                call_model,
+                CallModelInput(messages, streams.ActivityStream.from_handle(stream)),
+                start_to_close_timeout=timedelta(minutes=5),
+            )
+            if not reply.tool_calls:
+                await turns.publish(Turn("done", reply.text))
+                return reply.text
+            for call in reply.tool_calls:
+                messages.append(
+                    await workflow.execute_activity(
+                        run_tool, call, start_to_close_timeout=timedelta(minutes=5)
+                    )
+                )
+
+
+async def main(
+    client: Client, task_queue: str
+) -> list[streams.Record | streams.Closed]:
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[AgentTurns],
+        activities=[call_model, run_tool],
+    ):
+        handle = await client.start_workflow(
+            AgentTurns.run,
+            ["find it"],
+            id=f"agent-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+        await handle.result()
+        stream = streams.StreamClient(client).workflow_stream(handle, Events)
+        return [
+            item
+            async for item in stream.subscribe()
+            if isinstance(item, (streams.Record, streams.Closed))
+        ]

tests/contrib/streams/examples/client_creates_stream.py

@@ -0,0 +1,64 @@
+"""Brook API proposal, "A client creates a Stream and passes it to a Workflow".
+
+The client creates the Stream with its settings, passes it to the Workflow as data, reads it,
+and closes it, since the creator owns it.
+"""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from datetime import timedelta
+
+from temporalio import workflow
+from temporalio.client import Client
+from temporalio.contrib import streams
+from temporalio.contrib.streams import workflow as stream_workflow
+from temporalio.worker import Worker
+
+
+@dataclass
+class Turn:
+    kind: str
+    text: str = ""
+
+
+@streams.defn(name="events")
+class Events:
+    turns = streams.topic(Turn)
+
+
+@workflow.defn
+class Chat:
+    @workflow.run
+    async def run(self, stream: stream_workflow.WorkflowStream[Events]) -> None:
+        # The Workflow only writes. The client, or auto-close, closes the Stream.
+        await stream.topic(Events.turns).publish(
+            Turn("model", "hi"), Turn("last", "bye")
+        )
+
+
+async def main(client: Client, task_queue: str) -> list[Turn]:
+    async with Worker(client, task_queue=task_queue, workflows=[Chat]):
+        chat_stream = await streams.StreamClient(client).create(
+            Events,
+            id=f"chat-{uuid.uuid4()}",
+            retention=timedelta(days=1),
+            auto_close=timedelta(hours=1),
+        )
+        await client.start_workflow(
+            Chat.run,
+            # A Stream is data, so the Workflow gets it as an argument.
+            stream_workflow.WorkflowStream.from_handle(chat_stream),
+            id=f"chat-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+        turns = []
+        async for item in chat_stream.subscribe():
+            if not isinstance(item, streams.Record):
+                continue
+            turns.append(item.data)
+            if item.data.kind == "last":
+                break
+        await chat_stream.close()
+        return turns

tests/contrib/streams/examples/resume_from_offset.py

@@ -0,0 +1,87 @@
+"""Brook API proposal, "A client resumes a scoped Stream from an offset".
+
+The Workflow writes one Stream per turn with a ``name=`` override. The Reader saves the offset
+after each Record itself, stops, and later resumes from it with ``from_offset``.
+"""
+
+from __future__ import annotations
+
+import uuid
+from datetime import timedelta
+
+from temporalio import workflow
+from temporalio.client import Client, WorkflowHandle
+from temporalio.contrib import streams
+from temporalio.contrib.streams import workflow as stream_workflow
+from temporalio.worker import Worker
+
+
+@streams.defn(name="events")
+class Events:
+    tokens = streams.topic(str)
+
+
+@workflow.defn
+class Turns:
+    def __init__(self) -> None:
+        self.more: list[str] = []
+        self.done = False
+
+    @workflow.run
+    async def run(self, first: list[str]) -> None:
+        turn = await stream_workflow.create_stream(Events, name="turn-3")
+        tokens = turn.topic(Events.tokens)
+        await tokens.publish(*first)
+        await workflow.wait_condition(lambda: bool(self.more) or self.done)
+        await tokens.publish(*self.more)
+
+    @workflow.signal
+    def add(self, words: list[str]) -> None:
+        self.more = words
+
+
+class OffsetStore:
+    """Where the Reader keeps its place, such as a database row."""
+
+    def __init__(self) -> None:
+        self.offsets: dict[str, int] = {}
+
+    def load(self, key: str) -> int | None:
+        return self.offsets.get(key)
+
+    def save(self, key: str, offset: int) -> None:
+        self.offsets[key] = offset
+
+
+async def read_some(
+    client: Client, handle: WorkflowHandle[Turns, None], store: OffsetStore, count: int
+) -> list[str]:
+    turn3 = streams.StreamClient(client).workflow_stream(handle, Events, name="turn-3")
+    tokens = turn3.topic(Events.tokens)
+    seen = []
+    async for item in tokens.subscribe(
+        from_offset=store.load("turn-3"), create_wait=timedelta(seconds=10)
+    ):
+        if isinstance(item, streams.Closed):
+            break
+        if isinstance(item, streams.Superseded):
+            continue
+        seen.append(item.data)
+        # from_offset is inclusive, so the next read starts after this Record.
+        store.save("turn-3", item.offset + 1)
+        if len(seen) == count:
+            break
+    return seen
+
+
+async def main(client: Client, task_queue: str) -> tuple[list[str], list[str]]:
+    async with Worker(client, task_queue=task_queue, workflows=[Turns]):
+        handle = await client.start_workflow(
+            Turns.run, ["a", "b"], id=f"turns-{uuid.uuid4()}", task_queue=task_queue
+        )
+        store = OffsetStore()
+        before = await read_some(client, handle, store, 2)
+        await handle.signal(Turns.add, ["c", "d"])
+        after = await read_some(client, handle, store, 2)
+        await handle.result()
+        return before, after