[Brook S9] Added the Stream examples and the end-to-end Redis tests.

.github/CODEOWNERS

@@ -12,3 +12,5 @@
 # manage repo-wide concerns.
 /temporalio/contrib/workflow_streams/       @temporalio/ai-sdk     @temporalio/sdk
 /tests/contrib/workflow_streams/            @temporalio/ai-sdk     @temporalio/sdk
+/temporalio/contrib/stream/                 @temporalio/ai-sdk     @temporalio/sdk
+/tests/contrib/streams/                     @temporalio/ai-sdk     @temporalio/sdk

changelog/added/babbling-brook.md

@@ -0,0 +1,2 @@
+Add experimental Streams to `temporalio.common`. A Stream is an ordered log of Records that Workflows, Activities and clients write, and that clients and Activities read. Register a backend with `stream_backends=` on `Client.connect`, `Replayer` or `WorkflowEnvironment`, then use `workflow.create_stream`, `Client.create_stream` or `Client.stream`.
+Add a Redis Stream backend in `temporalio.contrib.stream`, and an in-memory one for tests in `temporalio.contrib.stream.testing`. The Redis backend needs a build with the bridge's `streams` feature, which the official builds have.

temporalio/common/__init__.py

@@ -1,4 +1,25 @@
-"""Common code used in the Temporal SDK."""
+"""Common code used in the Temporal SDK.
+
+This module also has the Streams API. A Stream is an ordered log of Records that Workflows,
+Activities and clients write and read. A Stream definition is a subclass of
+:py:class:`StreamDefinition` with ``@common.defn``. It gives the Stream name and its one payload
+type. Code creates a Stream before its first publish, and a create is create-or-get. Workflow
+code creates one with :py:func:`temporalio.workflow.create_stream`, and a client with
+:py:meth:`temporalio.client.Client.create_stream`. A client binds any Stream with
+:py:meth:`temporalio.client.Client.stream`. A Stream is data. A Workflow can pass it as an
+argument to an Activity, a child Workflow or the next run, or hand it out through a query. The
+parameter's type sets the handle class for that context, such as :py:class:`ActivityStream`.
+
+A :py:class:`StreamBackend` makes the Stream calls and registers a backend name. Examples are
+:py:class:`temporalio.contrib.stream.RedisStreamBackend` (``redis``) and
+:py:class:`temporalio.contrib.stream.testing.MemoryStreamBackend` (``memory``). Register
+backends with ``Client.connect(..., stream_backends=[...])``. A handle stores the name of its
+backend. With more than one registered backend, exactly one must have ``default=True``. A
+create with no ``backend=`` uses that backend. A Worker serves only one backend.
+
+.. warning::
+    The Streams API is experimental.
+"""
 
 from __future__ import annotations
 

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 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) |
+| Worker setup | [`worker_setup.py`](worker_setup.py) |
+
+The proposal's Nexus examples aren't in this directory.

tests/contrib/streams/examples/activity_streams_tokens.py

@@ -0,0 +1,98 @@
+"""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, common, workflow
+from temporalio.client import Client
+from temporalio.worker import Worker
+
+
+@common.defn(name="tokens")
+class Tokens(common.StreamDefinition[str]): ...
+
+
+@dataclass
+class CallModelInput:
+    prompt: str
+    tokens: common.ActivityStream[str]  # the Stream that the Workflow passes
+
+
+@dataclass
+class SummarizeInput:
+    text: str
+    tokens: common.SyncActivityStream[str]
+
+
+@activity.defn
+async def call_model(input: CallModelInput) -> int:
+    words = input.prompt.split()
+    # Each awaited publish has written its Record when it returns.
+    last = await input.tokens.publish(words[0])
+    # A fast model sends many small tokens, so the helper batches them.
+    async with common.BufferedPublisher(
+        input.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.tokens.publish(f"summary: {input.text}")
+
+
+@workflow.defn
+class Model:
+    @workflow.query
+    def stream(self) -> common.WorkflowStream[str]:
+        return workflow.stream_handle(Tokens)
+
+    @workflow.run
+    async def run(self, prompt: str) -> list[int]:
+        tokens = await workflow.create_stream(Tokens)
+        first = await workflow.execute_activity(
+            call_model,
+            CallModelInput(prompt, common.ActivityStream.from_handle(tokens)),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        summary = await workflow.execute_activity(
+            summarize,
+            SummarizeInput(prompt, common.SyncActivityStream.from_handle(tokens)),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        await tokens.close()
+        return [first, summary]
+
+
+async def main(
+    client: Client, task_queue: str
+) -> tuple[list[int], list[common.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 = await handle.query("stream", result_type=common.ClientStream[str])
+    records = [r async for r in stream.subscribe() if isinstance(r, common.Record)]
+    return offsets, records

tests/contrib/streams/examples/agent_publishes_turns.py

@@ -0,0 +1,116 @@
+"""Brook API proposal, "An agent Workflow publishes its turns".
+
+The Workflow creates its two Streams, turns and token deltas, and hands out the turns Stream
+through a ``stream`` query. It publishes each turn with an awaited ``publish``, and passes the
+tokens Stream to the model Activity in its one input object. A create is create-or-get, so the
+run after continue-as-new gets the same Streams. It closes both Streams before it returns, since
+a run's end doesn't close a Stream.
+"""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from datetime import timedelta
+
+from temporalio import activity, common, workflow
+from temporalio.client import Client
+from temporalio.worker import Worker
+
+
+@dataclass
+class Turn:
+    kind: str
+    text: str = ""
+
+
+@common.defn(name="turns")
+class Turns(common.StreamDefinition[Turn]): ...
+
+
+@common.defn(name="tokens")
+class Tokens(common.StreamDefinition[str]): ...
+
+
+@dataclass
+class Reply:
+    text: str
+    tool_calls: list[str] = field(default_factory=list)
+
+
+@dataclass
+class CallModelInput:
+    messages: list[str]
+    tokens: common.ActivityStream[str]
+
+
+@activity.defn
+async def call_model(input: CallModelInput) -> Reply:
+    await input.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 Agent:
+    @workflow.query
+    def stream(self) -> common.WorkflowStream[Turn]:
+        # The handle makes no call, so a reader can get it before the create.
+        return workflow.stream_handle(Turns)
+
+    @workflow.run
+    async def run(self, messages: list[str]) -> str:
+        # Create-or-get, so a retry or the run after continue-as-new gets the same Streams.
+        turns = await workflow.create_stream(Turns, retention=timedelta(days=1))
+        tokens = await workflow.create_stream(Tokens, retention=timedelta(days=1))
+        while True:
+            if workflow.info().is_continue_as_new_suggested():
+                workflow.continue_as_new(messages)
+            # The publish is one Local Activity. History keeps its offset, not the Turn.
+            await turns.publish(Turn("model"))
+            reply = await workflow.execute_activity(
+                call_model,
+                CallModelInput(messages, common.ActivityStream.from_handle(tokens)),
+                start_to_close_timeout=timedelta(minutes=5),
+            )
+            if not reply.tool_calls:
+                await turns.publish(Turn("done", reply.text))
+                # A run's end doesn't close the Streams, so the Workflow closes them.
+                await turns.close(result=reply.text)
+                await tokens.close()
+                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[common.Record | common.Closed]:
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[Agent],
+        activities=[call_model, run_tool],
+    ):
+        handle = await client.start_workflow(
+            Agent.run,
+            ["find it"],
+            id=f"agent-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+        await handle.result()
+        # The Workflow hands out its turns Stream, and the client reads it as a ClientStream.
+        stream = await handle.query("stream", result_type=common.ClientStream[Turn])
+        return [
+            item
+            async for item in stream.subscribe()
+            if isinstance(item, (common.Record, common.Closed))
+        ]

tests/contrib/streams/examples/client_creates_stream.py

@@ -0,0 +1,56 @@
+"""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. A create is create-or-get by Stream id.
+"""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from datetime import timedelta
+
+from temporalio import common, workflow
+from temporalio.client import Client
+from temporalio.worker import Worker
+
+
+@dataclass
+class Turn:
+    kind: str
+    text: str = ""
+
+
+@common.defn(name="turns")
+class Turns(common.StreamDefinition[Turn]): ...
+
+
+@workflow.defn
+class Chat:
+    @workflow.run
+    async def run(self, stream: common.WorkflowStream[Turn]) -> None:
+        # The Workflow only writes. The client closes the Stream.
+        await stream.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]):
+        turns = []
+        # The block closes the Stream at its end, or with the error that ends it.
+        async with client.create_stream(
+            Turns, id=f"chat-{uuid.uuid4()}", retention=timedelta(days=1)
+        ) as stream:
+            await client.start_workflow(
+                Chat.run,
+                # A Stream is data, so the Workflow gets it as an argument.
+                common.WorkflowStream.from_handle(stream),
+                id=f"chat-{uuid.uuid4()}",
+                task_queue=task_queue,
+            )
+            async for r in stream.subscribe():
+                if not isinstance(r, common.Record):
+                    continue
+                turns.append(r.data)
+                if r.data.kind == "last":
+                    break
+        return turns