From 712de6b6b1ebb298f3375f0d30ccb14b52c913f8 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 23:18:23 -0700 Subject: [PATCH 1/2] Added full examples for the streams API proposal. --- tests/contrib/streams/examples/README.md | 15 +++ tests/contrib/streams/examples/__init__.py | 0 .../examples/activity_streams_tokens.py | 99 +++++++++++++++ .../streams/examples/agent_publishes_turns.py | 107 ++++++++++++++++ .../streams/examples/client_creates_stream.py | 64 ++++++++++ .../streams/examples/resume_from_offset.py | 89 ++++++++++++++ .../contrib/streams/examples/test_examples.py | 91 ++++++++++++++ .../streams/examples/turn_scoped_stream.py | 114 ++++++++++++++++++ .../streams/examples/ui_backend_reads.py | 99 +++++++++++++++ .../contrib/streams/examples/worker_setup.py | 77 ++++++++++++ 10 files changed, 755 insertions(+) create mode 100644 tests/contrib/streams/examples/README.md create mode 100644 tests/contrib/streams/examples/__init__.py create mode 100644 tests/contrib/streams/examples/activity_streams_tokens.py create mode 100644 tests/contrib/streams/examples/agent_publishes_turns.py create mode 100644 tests/contrib/streams/examples/client_creates_stream.py create mode 100644 tests/contrib/streams/examples/resume_from_offset.py create mode 100644 tests/contrib/streams/examples/test_examples.py create mode 100644 tests/contrib/streams/examples/turn_scoped_stream.py create mode 100644 tests/contrib/streams/examples/ui_backend_reads.py create mode 100644 tests/contrib/streams/examples/worker_setup.py diff --git a/tests/contrib/streams/examples/README.md b/tests/contrib/streams/examples/README.md new file mode 100644 index 0000000000..9564002ee9 --- /dev/null +++ b/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. diff --git a/tests/contrib/streams/examples/__init__.py b/tests/contrib/streams/examples/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/contrib/streams/examples/activity_streams_tokens.py b/tests/contrib/streams/examples/activity_streams_tokens.py new file mode 100644 index 0000000000..49303c502a --- /dev/null +++ b/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 diff --git a/tests/contrib/streams/examples/agent_publishes_turns.py b/tests/contrib/streams/examples/agent_publishes_turns.py new file mode 100644 index 0000000000..5eff8daa4e --- /dev/null +++ b/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)) + ] diff --git a/tests/contrib/streams/examples/client_creates_stream.py b/tests/contrib/streams/examples/client_creates_stream.py new file mode 100644 index 0000000000..0bb9d5a810 --- /dev/null +++ b/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 diff --git a/tests/contrib/streams/examples/resume_from_offset.py b/tests/contrib/streams/examples/resume_from_offset.py new file mode 100644 index 0000000000..aa5196e33f --- /dev/null +++ b/tests/contrib/streams/examples/resume_from_offset.py @@ -0,0 +1,89 @@ +"""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 +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( + stream: streams.ClientStream[Events], store: OffsetStore, count: int +) -> list[str]: + seen = [] + tokens = stream.topic(Events.tokens) + 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 + ) + stream = streams.StreamClient(client).workflow_stream( + handle, Events, name="turn-3" + ) + store = OffsetStore() + before = await read_some(stream, store, 2) + await handle.signal(Turns.add, ["c", "d"]) + after = await read_some(stream, store, 2) + await handle.result() + return before, after diff --git a/tests/contrib/streams/examples/test_examples.py b/tests/contrib/streams/examples/test_examples.py new file mode 100644 index 0000000000..ec4939196e --- /dev/null +++ b/tests/contrib/streams/examples/test_examples.py @@ -0,0 +1,91 @@ +"""Runs each example end to end on the memory Store and on Redis, so none goes stale.""" + +import uuid + +from temporalio.client import Client +from temporalio.contrib import streams +from tests.contrib.streams.examples import ( + activity_streams_tokens, + agent_publishes_turns, + client_creates_stream, + resume_from_offset, + turn_scoped_stream, + ui_backend_reads, + worker_setup, +) + + +def task_queue() -> str: + return f"tq-{uuid.uuid4()}" + + +async def test_ui_backend_reads(stream_client: Client): + screen = await ui_backend_reads.main(stream_client, task_queue()) + assert [(record.topic, record.data) for record in screen.shown] == [ + ("turns", ui_backend_reads.Turn("model")), + ("tokens", "hello"), + ("tokens", "there"), + ("turns", ui_backend_reads.Turn("done", "HELLO THERE")), + ] + assert screen.result == "HELLO THERE" + + +async def test_resume_from_offset(stream_client: Client): + assert await resume_from_offset.main(stream_client, task_queue()) == ( + ["a", "b"], + ["c", "d"], + ) + + +async def test_client_creates_stream(stream_client: Client): + turns = await client_creates_stream.main(stream_client, task_queue()) + assert [turn.kind for turn in turns] == ["model", "last"] + + +async def test_agent_publishes_turns(stream_client: Client): + items = await agent_publishes_turns.main(stream_client, task_queue()) + data = [item.data for item in items if isinstance(item, streams.Record)] + Turn = agent_publishes_turns.Turn + assert data == [Turn("model"), "find", "it", Turn("model"), "search", "result"] + [ + Turn("done", "the answer") + ] + assert isinstance(items[-1], streams.Closed) and items[-1].result == "the answer" + + +async def test_activity_streams_tokens(stream_client: Client): + offsets, records = await activity_streams_tokens.main(stream_client, task_queue()) + assert [record.data for record in records] == [ + "one", + "two", + "three", + "summary: one two three", + ] + assert offsets == [records[0].offset, records[-1].offset] + + +async def test_turn_scoped_stream(stream_client: Client): + assert await turn_scoped_stream.main(stream_client, task_queue()) == { + 0: ["a", "b"], + 1: ["c"], + 2: ["d", "e", "f"], + } + + +async def test_worker_setup_on_memory(client: Client): + target = client.service_client.config.target_host + test_client = await worker_setup.connect_for_tests(target) + assert await worker_setup.main(test_client, task_queue()) == [ + "hello ada", + "bye ada", + ] + + +async def test_worker_setup_on_redis(client: Client, redis_url: str): + target = client.service_client.config.target_host + redis_client = await worker_setup.connect_with_redis(target, redis_url) + assert await worker_setup.main(redis_client, task_queue()) == [ + "hello ada", + "bye ada", + ] + migrating = await worker_setup.connect_while_migrating(target, redis_url, redis_url) + assert await worker_setup.create_on_second_backend(migrating) == "redis-b" diff --git a/tests/contrib/streams/examples/turn_scoped_stream.py b/tests/contrib/streams/examples/turn_scoped_stream.py new file mode 100644 index 0000000000..d615bcc63f --- /dev/null +++ b/tests/contrib/streams/examples/turn_scoped_stream.py @@ -0,0 +1,114 @@ +"""Brook API proposal, "A Workflow scopes a Stream to one turn and closes it". + +Each turn gets its own Stream, named with ``name=``, which ``async with`` closes at the end of +the turn. An index Stream lists the turns, and continue-as-new passes it to the next run. +""" + +from __future__ import annotations + +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 TurnStarted: + n: int + stream: streams.ClientStream[Events] + + +@streams.defn(name="index") +class TurnIndex: + started = streams.topic(TurnStarted) + + +@dataclass +class CallModelInput: + prompt: str + stream: streams.ActivityStream[Events] + + +@activity.defn +async def call_model_into(input: CallModelInput) -> None: + await input.stream.topic(Events.tokens).publish(*input.prompt.split()) + + +@dataclass +class ChatTurnsInput: + prompts: list[str] + # The first turn number of this run. A Stream name must be unique in the run chain. + next: int = 0 + # From the run before continue-as-new. + index: stream_workflow.WorkflowStream[TurnIndex] | None = None + + +@workflow.defn +class ChatTurns: + @workflow.run + async def run(self, input: ChatTurnsInput) -> None: + index = input.index or await stream_workflow.create_stream( + TurnIndex, retention=timedelta(days=1) + ) + started = index.topic(TurnIndex.started) + batch = input.prompts[:2] + for i, prompt in enumerate(batch): + async with await stream_workflow.create_stream( + Events, + name=f"turn-{input.next + i}", + retention=timedelta(days=1), + close_with_run=False, + ) as turn: + # A Stream is data, so the index Record carries the turn's Stream. + await started.publish( + TurnStarted(input.next + i, streams.ClientStream.from_handle(turn)) + ) + await workflow.execute_activity( + call_model_into, + CallModelInput(prompt, streams.ActivityStream.from_handle(turn)), + start_to_close_timeout=timedelta(minutes=5), + ) + if len(input.prompts) > len(batch): + workflow.continue_as_new( + ChatTurnsInput(input.prompts[2:], input.next + len(batch), index) + ) + + +async def main(client: Client, task_queue: str) -> dict[int, list[str]]: + async with Worker( + client, + task_queue=task_queue, + workflows=[ChatTurns], + activities=[call_model_into], + ): + handle = await client.start_workflow( + ChatTurns.run, + ChatTurnsInput(["a b", "c", "d e f"]), + id=f"chat-{uuid.uuid4()}", + task_queue=task_queue, + ) + await handle.result() + stream_client = streams.StreamClient(client) + index = stream_client.workflow_stream(handle, TurnIndex) + turns = {} + async for item in index.subscribe(): + if not isinstance(item, streams.Record): + continue + # A Stream from a payload has no client, so the reader gets it again by id. + turn = stream_client.get_handle(Events, id=item.data.stream.id) + turns[item.data.n] = [ + token.data + async for token in turn.subscribe() + if isinstance(token, streams.Record) + ] + return turns diff --git a/tests/contrib/streams/examples/ui_backend_reads.py b/tests/contrib/streams/examples/ui_backend_reads.py new file mode 100644 index 0000000000..f058e14949 --- /dev/null +++ b/tests/contrib/streams/examples/ui_backend_reads.py @@ -0,0 +1,99 @@ +"""Brook API proposal, "A UI backend reads a Workflow's Stream". + +The agent Workflow creates the ``events`` Stream. The UI backend starts it, gets the Stream +from the Workflow handle with no call to Temporal, and renders each item as it comes. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +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 CallModelInput: + prompt: str + stream: streams.ActivityStream[Events] + + +@activity.defn +async def call_model(input: CallModelInput) -> str: + tokens = input.stream.topic(Events.tokens) + for word in input.prompt.split(): + await tokens.publish(word) + return input.prompt.upper() + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + stream = await stream_workflow.create_stream(Events) + turns = stream.topic(Events.turns) + await turns.publish(Turn("model")) + reply = await workflow.execute_activity( + call_model, + CallModelInput(prompt, streams.ActivityStream.from_handle(stream)), + start_to_close_timeout=timedelta(minutes=5), + ) + await turns.publish(Turn("done", reply)) + return reply + + +class Screen: + """Stands in for the UI. It keeps what it shows, so a withdrawal can take items back.""" + + def __init__(self) -> None: + self.shown: list[streams.Record[Any]] = [] + self.result: Any = None + + def render(self, record: streams.Record[Any]) -> None: + self.shown.append(record) + + def unrender_where(self, superseded: streams.Superseded) -> None: + self.shown = [ + record for record in self.shown if not superseded.withdraws(record) + ] + + +async def main(client: Client, task_queue: str) -> Screen: + async with Worker( + client, task_queue=task_queue, workflows=[Agent], activities=[call_model] + ): + handle = await client.start_workflow( + Agent.run, "hello there", id=f"agent-{uuid.uuid4()}", task_queue=task_queue + ) + stream = streams.StreamClient(client).workflow_stream(handle, Events) + screen = Screen() + # The Workflow may not have created the Stream yet, so the read waits for it. + async for item in stream.subscribe( + Events.turns, Events.tokens, create_wait=timedelta(seconds=10) + ): + if isinstance(item, streams.Closed): + screen.result = item.result if item.failure is None else item.failure + elif isinstance(item, streams.Superseded): + screen.unrender_where(item) + else: + # item.data is a Turn or a str, by item.topic. + screen.render(item) + return screen diff --git a/tests/contrib/streams/examples/worker_setup.py b/tests/contrib/streams/examples/worker_setup.py new file mode 100644 index 0000000000..7cb346c43a --- /dev/null +++ b/tests/contrib/streams/examples/worker_setup.py @@ -0,0 +1,77 @@ +"""Brook API proposal, "Worker setup". + +The Stream plugin goes on the client, and the Worker takes the client's plugins. Each plugin +registers a backend name. Tests use ``MemoryStreamPlugin`` with no change to Workflow code. +Sync Workflow code, such as ``@workflow.init``, uses the ``_nowait`` calls. +""" + +from __future__ import annotations + +import uuid + +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.contrib.streams.redis import RedisStreamPlugin +from temporalio.contrib.streams.testing import MemoryStreamPlugin +from temporalio.worker import Worker + + +@streams.defn(name="events") +class Events: + tokens = streams.topic(str) + + +@workflow.defn +class Greeter: + @workflow.init + def __init__(self, name: str) -> None: + # The init can't await, so it creates and publishes with the _nowait calls. + self.stream = stream_workflow.create_stream_nowait(Events) + self.stream.topic(Events.tokens).publish_nowait(f"hello {name}") + + @workflow.run + async def run(self, name: str) -> str: + await self.stream.topic(Events.tokens).publish(f"bye {name}") + return name + + +async def connect_with_redis(target_host: str, redis_url: str) -> Client: + return await Client.connect(target_host, plugins=[RedisStreamPlugin(redis_url)]) + + +async def connect_for_tests(target_host: str) -> Client: + return await Client.connect(target_host, plugins=[MemoryStreamPlugin()]) + + +async def connect_while_migrating(target_host: str, url_a: str, url_b: str) -> Client: + # Two backends during a move. Creates with no backend= go to the default one. + return await Client.connect( + target_host, + plugins=[ + RedisStreamPlugin(url_a, name="redis-a", default=True), + RedisStreamPlugin(url_b, name="redis-b"), + ], + ) + + +async def main(client: Client, task_queue: str) -> list[str]: + async with Worker(client, task_queue=task_queue, workflows=[Greeter]): + handle = await client.start_workflow( + Greeter.run, "ada", id=f"greeter-{uuid.uuid4()}", task_queue=task_queue + ) + await handle.result() + stream = streams.StreamClient(client).workflow_stream(handle, Events) + return [ + item.data + async for item in stream.subscribe() + if isinstance(item, streams.Record) + ] + + +async def create_on_second_backend(client: Client) -> str | None: + stream = await streams.StreamClient(client).create( + Events, id=f"chat-{uuid.uuid4()}", backend="redis-b" + ) + return stream.backend From 0843a727ff0ec02e9ce09872b72ef40586c09500 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 23:53:46 -0700 Subject: [PATCH 2/2] Kept each example's proposal snippet in one block. --- .../contrib/streams/examples/resume_from_offset.py | 14 ++++++-------- tests/contrib/streams/examples/worker_setup.py | 7 ++++++- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/contrib/streams/examples/resume_from_offset.py b/tests/contrib/streams/examples/resume_from_offset.py index aa5196e33f..e0c3f81a08 100644 --- a/tests/contrib/streams/examples/resume_from_offset.py +++ b/tests/contrib/streams/examples/resume_from_offset.py @@ -10,7 +10,7 @@ from datetime import timedelta from temporalio import workflow -from temporalio.client import Client +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 @@ -54,10 +54,11 @@ def save(self, key: str, offset: int) -> None: async def read_some( - stream: streams.ClientStream[Events], store: OffsetStore, count: int + 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 = [] - tokens = stream.topic(Events.tokens) async for item in tokens.subscribe( from_offset=store.load("turn-3"), create_wait=timedelta(seconds=10) ): @@ -78,12 +79,9 @@ async def main(client: Client, task_queue: str) -> tuple[list[str], list[str]]: handle = await client.start_workflow( Turns.run, ["a", "b"], id=f"turns-{uuid.uuid4()}", task_queue=task_queue ) - stream = streams.StreamClient(client).workflow_stream( - handle, Events, name="turn-3" - ) store = OffsetStore() - before = await read_some(stream, store, 2) + before = await read_some(client, handle, store, 2) await handle.signal(Turns.add, ["c", "d"]) - after = await read_some(stream, store, 2) + after = await read_some(client, handle, store, 2) await handle.result() return before, after diff --git a/tests/contrib/streams/examples/worker_setup.py b/tests/contrib/streams/examples/worker_setup.py index 7cb346c43a..e29bda15f0 100644 --- a/tests/contrib/streams/examples/worker_setup.py +++ b/tests/contrib/streams/examples/worker_setup.py @@ -56,8 +56,13 @@ async def connect_while_migrating(target_host: str, url_a: str, url_b: str) -> C ) +def run_worker(client: Client, task_queue: str) -> Worker: + # The Worker takes the client's plugins, so it needs no Stream setup of its own. + return Worker(client, task_queue=task_queue, workflows=[Greeter]) + + async def main(client: Client, task_queue: str) -> list[str]: - async with Worker(client, task_queue=task_queue, workflows=[Greeter]): + async with run_worker(client, task_queue): handle = await client.start_workflow( Greeter.run, "ada", id=f"greeter-{uuid.uuid4()}", task_queue=task_queue )