[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. A Stream is an ordered log of Records that Workflows, Activities and clients write, and that clients and Activities read. `temporalio.common` has the shared types, such as the `StreamHandle` data, `temporalio.workflow.WorkflowStream` makes the calls in Workflow code, and `temporalio.client.Stream` makes them elsewhere. Register a backend with `stream_backends=` on `Client.connect`, `Replayer` or `WorkflowEnvironment`, then use `workflow.create_stream`, `Client.create_stream` or `Client.get_stream`, and `to_stream` for a handle that came as data.
+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,31 @@
-"""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`, which gives a
+:py:class:`temporalio.workflow.WorkflowStream`. A client creates one with
+:py:meth:`temporalio.client.Client.create_stream`, which gives a
+:py:class:`temporalio.client.Stream`, and gets any Stream by its id with
+:py:meth:`temporalio.client.Client.get_stream`.
+
+A Stream's :py:class:`StreamHandle` is data with no calls. Code passes it as an argument to a
+Workflow, an Activity, a child Workflow or the next run, or hands it out through a query.
+:py:func:`temporalio.workflow.to_stream` and :py:meth:`temporalio.client.Client.to_stream` give
+the Stream of a handle. In an Activity, a Stream from ``activity.client()`` writes as the
+attempt.
+
+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,16 @@
+# Examples for the 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 Stream by its id | [`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) |
+| A Workflow that gets a handle as data | [`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,115 @@
+"""Streams API proposal, "An Activity streams tokens to the Workflow's Stream".
+
+The Workflow passes its Stream's handle to the Activity as a field of the input. An async
+Activity gets the Stream from ``activity.client()``, so it writes as the attempt, and batches the
+model's deltas with ``BufferedPublisher``. A sync Activity has no ``activity.client()``, so it
+keeps the Worker's client. It wraps its Stream in ``SyncStream``, whose calls block its thread,
+so it works in a thread pool executor.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import uuid
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from datetime import timedelta
+
+from temporalio import activity, common, workflow
+from temporalio.client import BufferedPublisher, Client, SyncStream
+from temporalio.worker import Worker
+
+
+@common.defn(name="tokens")
+class Tokens(common.StreamDefinition[str]): ...
+
+
+class FakeResponse:
+    """A model's streamed response, one delta for each word."""
+
+    def __init__(self, messages: list[str]) -> None:
+        self.words = messages[-1].split()
+
+    async def __aiter__(self) -> AsyncIterator[str]:
+        for word in self.words:
+            yield word
+
+    def reply(self) -> str:
+        return " ".join(self.words).upper()
+
+
+@dataclass
+class CallModelInput:
+    messages: list[str]
+    tokens: common.StreamHandle[str]  # the handle that the Workflow passes
+
+
+@dataclass
+class SummarizeInput:
+    text: str
+    tokens: common.StreamHandle[str]
+
+
+@activity.defn
+async def call_model(input: CallModelInput) -> str:
+    stream = activity.client().to_stream(input.tokens)  # writes as this attempt
+    response = FakeResponse(input.messages)
+    async with BufferedPublisher(stream) as tokens:  # batches deltas
+        async for delta in response:
+            tokens.publish(delta)
+    return response.reply()
+
+
+class Summarizer:
+    def __init__(self, client: Client) -> None:
+        self._client = client  # the Worker's client
+
+    @activity.defn
+    def summarize(self, input: SummarizeInput) -> int:
+        tokens = SyncStream(self._client.to_stream(input.tokens))
+        # The sync publish blocks until the Record is written, and returns its offset.
+        return tokens.publish(f"summary: {input.text}")
+
+
+@workflow.defn
+class Model:
+    @workflow.query
+    def stream(self) -> common.StreamHandle[str]:
+        return workflow.get_stream(Tokens).handle
+
+    @workflow.run
+    async def run(self, prompt: str) -> int:
+        tokens = await workflow.create_stream(Tokens)
+        reply = await workflow.execute_activity(
+            call_model,
+            CallModelInput([prompt], tokens.handle),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        summary = await workflow.execute_activity_method(
+            Summarizer.summarize,
+            SummarizeInput(reply, tokens.handle),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        await tokens.close()
+        return summary
+
+
+async def main(client: Client, task_queue: str) -> tuple[int, list[common.Record[str]]]:
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        async with Worker(
+            client,
+            task_queue=task_queue,
+            workflows=[Model],
+            activities=[call_model, Summarizer(client).summarize],
+            activity_executor=executor,
+        ):
+            handle = await client.start_workflow(
+                Model.run,
+                "one two three",
+                id=f"model-{uuid.uuid4()}",
+                task_queue=task_queue,
+            )
+            summary = await handle.result()
+            stream = client.to_stream(await handle.query(Model.stream))
+    records = [r async for r in stream.subscribe() if isinstance(r, common.Record)]
+    return summary, records

tests/contrib/streams/examples/agent_publishes_turns.py

@@ -0,0 +1,125 @@
+"""Streams API proposal, "An agent Workflow publishes its turns".
+
+The Workflow creates its two Streams, turns and token deltas, and hands out the turns Stream's
+handle through a ``stream`` query. It publishes each turn with an awaited ``publish``, and passes
+the tokens handle 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 BufferedPublisher, Client
+from temporalio.worker import Worker
+
+MINUTE = timedelta(minutes=1)
+
+
+@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.StreamHandle[str]  # the handle that the Workflow passes
+
+
+@activity.defn
+async def call_model(input: CallModelInput) -> Reply:
+    stream = activity.client().to_stream(input.tokens)  # writes as this attempt
+    async with BufferedPublisher(stream) as tokens:  # batches deltas
+        for delta in input.messages[-1].split():
+            tokens.publish(delta)
+    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:
+    def __init__(self) -> None:
+        # get_stream makes no call, so the query has the handle before the create.
+        self._turns = workflow.get_stream(Turns)
+
+    @workflow.query
+    def stream(self) -> common.StreamHandle[Turn]:  # one way to hand it out
+        return self._turns.handle
+
+    @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.
+        self._turns = turns = await workflow.create_stream(
+            Turns
+        )  # the default Retention
+        # Token deltas matter only to live readers, so they get a short Retention.
+        tokens = await workflow.create_stream(Tokens, retention=timedelta(hours=1))
+        while True:
+            if workflow.info().is_continue_as_new_suggested():
+                workflow.continue_as_new(messages)
+            await turns.publish(Turn("model"))  # one Local Activity
+            reply = await workflow.execute_activity(
+                call_model,
+                CallModelInput(messages, tokens.handle),
+                start_to_close_timeout=MINUTE,
+            )
+            if not reply.tool_calls:
+                await turns.publish(Turn("done", reply.text))
+                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=MINUTE
+                    )
+                )
+
+
+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 handle, and the query's client reads the Stream.
+        stream = client.to_stream(await handle.query(Agent.stream), payload_type=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,74 @@
+"""Streams API proposal, "A client creates a Stream and passes it to a Workflow".
+
+The client creates the Stream with its settings, passes its handle to the Workflow as data,
+reads it, and closes it, since the creator owns it. A create is create-or-get by Stream id. The
+Workflow gets the handle as data, turns it into a Stream to write, and passes the same handle to
+an Activity.
+"""
+
+from __future__ import annotations
+
+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
+
+
+@dataclass
+class Turn:
+    kind: str
+    text: str = ""
+
+
+@common.defn(name="turns")
+class Turns(common.StreamDefinition[Turn]): ...
+
+
+SECONDS_30 = timedelta(seconds=30)
+
+
+@activity.defn
+async def notify_user(handle: common.StreamHandle[Turn]) -> str:
+    # A real Activity sends the Stream id to the user's device, which reads it by id.
+    return handle.id
+
+
+@workflow.defn
+class Chat:
+    @workflow.run
+    async def run(self, handle: common.StreamHandle[Turn]) -> None:
+        # The Workflow only writes. The client closes the Stream.
+        turns = workflow.to_stream(handle)
+        turns.publish(Turn("model", "hi"))
+        await workflow.execute_activity(
+            notify_user, handle, start_to_close_timeout=SECONDS_30
+        )
+        await turns.publish(Turn("last", "bye"))
+
+
+def is_last_turn(record: common.Record[Turn]) -> bool:
+    return record.data.kind == "last"
+
+
+async def main(client: Client, task_queue: str) -> list[Turn]:
+    async with Worker(
+        client, task_queue=task_queue, workflows=[Chat], activities=[notify_user]
+    ):
+        chat_id = f"chat-{uuid.uuid4()}"
+        turns = []
+        async with client.create_stream(  # the block end closes it
+            Turns, id=chat_id, retention=timedelta(days=1)
+        ) as stream:
+            await client.start_workflow(  # the handle is data
+                Chat.run, stream.handle, id=chat_id, task_queue=task_queue
+            )
+            async for r in stream.subscribe():
+                if not isinstance(r, common.Record):
+                    continue
+                turns.append(r.data)
+                if is_last_turn(r):
+                    break
+        return turns