Added the Stream examples and the end-to-end Redis tests.

Author: moetemp

.github/CODEOWNERS

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

changelog/added/meandering-river.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.streams` has the shared types, such as the `StreamRef` 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[T]`, `Client.create_stream[T]` or `Client.get_stream[T]`, and `to_stream` for a ref 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/streams/__init__.py

@@ -1,5 +1,34 @@
 """The Streams API, which Workflows, Activities and clients share.
 
+A Stream is an ordered log of Records that Workflows, Activities and clients write, and that
+clients and Activities read. A Stream has one payload type, and a call takes it as a type
+argument, as in ``client.get_stream[Turn]("chat-42")``. With no type argument, the payload type
+is ``Any``, and a read gives JSON values.
+
+Code creates a Stream before its first publish. :py:data:`temporalio.workflow.create_stream` and
+:py:attr:`temporalio.client.Client.create_stream` fail with :py:class:`StreamAlreadyExistsError` if
+the id exists, and ``create_or_get_stream`` gets the Stream instead. In Workflow code, each run of a
+chain sends the same create for an id, so a retry run or a cron run gets the Stream that an earlier
+run made. Workflow code gets a :py:class:`temporalio.workflow.WorkflowStream`, and other code a
+:py:class:`temporalio.client.Stream`. A client gets any Stream by its id with
+:py:attr:`temporalio.client.Client.get_stream`, with no call.
+
+A Stream's :py:class:`StreamRef` 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 ref. In an Activity, a Stream from ``activity.client()`` writes as the attempt.
+
+A read gives :py:class:`Record` and :py:class:`Superseded` items, and the ``kind`` of each item
+narrows its type. A close carries no value. A read of a closed Stream ends once it gives the last
+item.
+
+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 ref keeps 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.
 """

tests/contrib/streams/examples/README.md

@@ -0,0 +1,15 @@
+# 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 UI backend reads a Stream by its id, its resume notes | [`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/__init__.py

@@ -0,0 +1,137 @@
+"""Streams API proposal, "An Activity streams tokens to the Workflow's Stream".
+
+The Workflow passes its Stream's ref 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, streams, workflow
+from temporalio.client import BufferedPublisher, Client, SyncStream
+from temporalio.worker import Worker
+
+Message = str
+
+
+@dataclass
+class Reply:
+    text: str
+
+
+class FakeResponse:
+    """A model's streamed response, one delta for each word."""
+
+    def __init__(self, messages: list[Message]) -> None:
+        self.words = messages[-1].split()
+
+    async def __aiter__(self) -> AsyncIterator[str]:
+        for word in self.words:
+            yield word
+
+    def reply(self) -> Reply:
+        return Reply(" ".join(self.words).upper())
+
+
+class FakeLLM:
+    """Stands in for a model client."""
+
+    def stream(self, messages: list[Message]) -> FakeResponse:
+        return FakeResponse(messages)
+
+
+llm = FakeLLM()
+
+
+@dataclass
+class CallModelInput:
+    messages: list[Message]
+    tokens: streams.StreamRef[str]  # the ref that the Workflow passes
+
+
+@dataclass
+class SummarizeInput:
+    text: str
+    tokens: streams.StreamRef[str]
+
+
+@activity.defn
+async def call_model(input: CallModelInput) -> Reply:
+    stream = activity.client().to_stream(input.tokens)  # writes as this attempt
+    response = llm.stream(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:
+    def __init__(self) -> None:
+        self._tokens: workflow.WorkflowStream[str] | None = None
+
+    @workflow.query
+    def stream(self) -> streams.StreamRef[str] | None:
+        return self._tokens.ref if self._tokens else None
+
+    @workflow.run
+    async def run(self, prompt: str) -> int:
+        self._tokens = tokens = await workflow.create_stream[str](
+            f"{workflow.info().workflow_id}-tokens"
+        )
+        reply = await workflow.execute_activity(
+            call_model,
+            CallModelInput([prompt], tokens.ref),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        summary = await workflow.execute_activity_method(
+            Summarizer.summarize,
+            SummarizeInput(reply.text, tokens.ref),
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+        await tokens.close()
+        return summary
+
+
+async def main(
+    client: Client, task_queue: str
+) -> tuple[int, list[streams.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()
+            ref = await handle.query(Model.stream)
+    assert ref is not None
+    stream = client.to_stream(ref)
+    records = [r async for r in stream.subscribe() if isinstance(r, streams.Record)]
+    return summary, records

tests/contrib/streams/examples/activity_streams_tokens.py

@@ -0,0 +1,136 @@
+"""Streams API proposal, "An agent Workflow publishes its turns".
+
+The Workflow creates its two Streams, turns and token deltas, under ids of its choice, and hands
+out the turns Stream's ref through a ``stream`` query. It publishes each turn with an awaited
+``publish``, and passes the tokens ref to the model Activity in its one input object. A run
+after continue-as-new gets the earlier run's refs and creates nothing. 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, streams, workflow
+from temporalio.client import BufferedPublisher, Client
+from temporalio.worker import Worker
+
+MINUTE = timedelta(minutes=1)
+
+Message = str
+
+
+@dataclass
+class Turn:
+    kind: str
+    text: str | None = None
+
+
+@dataclass
+class Reply:
+    text: str
+    tool_calls: list[str] = field(default_factory=list)
+
+
+@dataclass
+class CallModelInput:
+    messages: list[Message]
+    tokens: streams.StreamRef[str]  # the ref 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) -> Message:
+    return f"{call} result"
+
+
+@dataclass
+class AgentInput:
+    session_id: str
+    messages: list[Message]
+    # The run before Continue-as-New sets the refs.
+    turns: streams.StreamRef[Turn] | None = None
+    tokens: streams.StreamRef[str] | None = None
+
+
+@workflow.defn
+class Agent:
+    def __init__(self) -> None:
+        self._turns: workflow.WorkflowStream[Turn] | None = None  # run creates it
+
+    @workflow.query
+    def stream(self) -> streams.StreamRef[Turn] | None:  # one way to hand it out
+        return self._turns.ref if self._turns else None  # None before the create
+
+    @workflow.run
+    async def run(self, input: AgentInput) -> str:
+        messages = input.messages
+        # A Workflow creates only standalone Streams, under an id of its choice. A run after
+        # Continue-as-New gets the earlier run's refs and creates nothing.
+        if input.turns and input.tokens:
+            self._turns = turns = workflow.to_stream(input.turns)
+            tokens = workflow.to_stream(input.tokens)
+        else:
+            self._turns = turns = await workflow.create_stream[Turn](
+                f"{input.session_id}-turns"
+            )  # the default Retention
+            # Token deltas matter only to live readers, so they get a short Retention.
+            tokens = await workflow.create_stream[str](
+                f"{input.session_id}-tokens", retention=timedelta(hours=1)
+            )
+        while True:
+            if workflow.info().is_continue_as_new_suggested():
+                workflow.continue_as_new(  # forward the refs, so the next run keeps the Streams
+                    AgentInput(input.session_id, messages, turns.ref, tokens.ref)
+                )
+            await turns.publish(Turn("model"))  # one Local Activity
+            reply = await workflow.execute_activity(
+                call_model,
+                CallModelInput(messages, tokens.ref),
+                start_to_close_timeout=MINUTE,
+            )
+            if not reply.tool_calls:
+                await turns.publish(Turn("done", reply.text))
+                await turns.close()  # the done Turn carries the reply
+                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[streams.Record[Turn]]:
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[Agent],
+        activities=[call_model, run_tool],
+    ):
+        session_id = f"session-{uuid.uuid4()}"
+        handle = await client.start_workflow(
+            Agent.run,
+            AgentInput(session_id, ["find it"]),
+            id=session_id,
+            task_queue=task_queue,
+        )
+        await handle.result()
+        # The Workflow hands out its turns ref, and the query's client reads the Stream.
+        ref = await handle.query(Agent.stream)
+        assert ref is not None
+        stream = client.to_stream(ref, payload_type=Turn)
+        return [r async for r in stream.subscribe() if isinstance(r, streams.Record)]