Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .github/CODEOWNERS
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 2 additions & 0 deletions changelog/added/meandering-river.md
Original file line number Diff line number Diff line change
@@ -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.
29 changes: 29 additions & 0 deletions temporalio/streams/__init__.py
Original file line number Diff line number Diff line change
@@ -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.
"""
Expand Down
15 changes: 15 additions & 0 deletions tests/contrib/streams/examples/README.md
Original file line number Diff line number Diff line change
@@ -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.
Empty file.
137 changes: 137 additions & 0 deletions tests/contrib/streams/examples/activity_streams_tokens.py
Original file line number Diff line number Diff line change
@@ -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
136 changes: 136 additions & 0 deletions tests/contrib/streams/examples/agent_publishes_turns.py
Original file line number Diff line number Diff line change
@@ -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)]
Loading
Loading